"""Title capture callback for runs."""

from __future__ import annotations

from typing import Any
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler


class RunTitleCallback(BaseCallbackHandler):
    """Capture title generated by title middleware LLM calls or custom events."""

    def __init__(self) -> None:
        super().__init__()
        self._title: str | None = None

    def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None:
        if self._identify_caller(kwargs) != "middleware:title":
            return
        try:
            message = response.generations[0][0].message
        except (IndexError, AttributeError):
            return
        content = getattr(message, "content", "")
        if isinstance(content, str) and content:
            self._title = content.strip().strip('"').strip("'")[:200]

    def on_custom_event(self, name: str, data: Any, *, run_id: UUID, **kwargs: Any) -> None:
        if name not in {"title", "thread_title", "middleware:title"}:
            return
        if isinstance(data, str):
            self._title = data.strip()[:200]
            return
        if isinstance(data, dict):
            title = data.get("title")
            if isinstance(title, str):
                self._title = title.strip()[:200]

    def title(self) -> str | None:
        return self._title

    def _identify_caller(self, kwargs: dict[str, Any]) -> str:
        for tag in kwargs.get("tags") or []:
            if isinstance(tag, str) and (
                tag.startswith("subagent:")
                or tag.startswith("middleware:")
                or tag == "lead_agent"
            ):
                return tag
        return "lead_agent"