"""Middleware for automatic thread title generation."""

import copy
import logging
import re
from pathlib import Path
from typing import TYPE_CHECKING, Any, NotRequired, override

from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langgraph.runtime import Runtime

from deerflow.config.title_config import get_title_config
from deerflow.models import create_chat_model
from deerflow.utils.time import now_iso

if TYPE_CHECKING:
    from deerflow.config.app_config import AppConfig
    from deerflow.config.title_config import TitleConfig

logger = logging.getLogger(__name__)
_TITLE_DEBUG_LOGGER_NAME = "title_generation_debug"


def _get_title_debug_logger() -> logging.Logger:
    debug_logger = logging.getLogger(_TITLE_DEBUG_LOGGER_NAME)
    if getattr(debug_logger, "_title_file_configured", False):
        return debug_logger

    log_path = Path("logs/title_generation_debug.log")
    if not log_path.is_absolute():
        log_path = Path.cwd() / log_path
    log_path.parent.mkdir(parents=True, exist_ok=True)

    handler = logging.FileHandler(log_path, encoding="utf-8")
    handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
    debug_logger.addHandler(handler)
    debug_logger.setLevel(logging.INFO)
    debug_logger.propagate = False
    setattr(debug_logger, "_title_file_configured", True)
    return debug_logger


def log_title_debug(message: str, *args: Any) -> None:
    logger.warning(message, *args)
    _get_title_debug_logger().info(message, *args)


class TitleMiddlewareState(AgentState):
    """Compatible with the `ThreadState` schema."""

    title: NotRequired[str | None]
    title_provisional: NotRequired[bool]


class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
    """Automatically generate a title for the thread after the first user message."""

    state_schema = TitleMiddlewareState

    def __init__(self, *, app_config: "AppConfig | None" = None, title_config: "TitleConfig | None" = None):
        super().__init__()
        self._app_config = app_config
        self._title_config = title_config

    def _get_title_config(self):
        if self._title_config is not None:
            return self._title_config
        if self._app_config is not None:
            return self._app_config.title
        return get_title_config()

    def _normalize_content(self, content: object) -> str:
        if isinstance(content, str):
            return content

        if isinstance(content, list):
            parts = [self._normalize_content(item) for item in content]
            return "\n".join(part for part in parts if part)

        if isinstance(content, dict):
            text_value = content.get("text")
            if isinstance(text_value, str):
                return text_value

            nested_content = content.get("content")
            if nested_content is not None:
                return self._normalize_content(nested_content)

        return ""

    def _message_type(self, message: object) -> str:
        if isinstance(message, dict):
            return str(message.get("type") or message.get("role") or "").lower()
        return str(getattr(message, "type", "") or getattr(message, "role", "") or "").lower()

    def _message_content(self, message: object) -> object:
        if isinstance(message, dict):
            return message.get("content", "")
        return getattr(message, "content", "")

    def _should_set_provisional_title(self, state: TitleMiddlewareState) -> bool:
        """Check if we should set an instant, question-derived placeholder title.

        Runs in ``before_model`` so the thread shows a sensible title the moment
        the user sends their first message — eliminating the "untitled" window
        while the (possibly long) first answer streams. No LLM call, no latency.
        """
        config = self._get_title_config()
        if not config.enabled:
            return False

        # Already has a title (provisional or final) — nothing to do.
        if state.get("title"):
            return False

        # Only on the very first turn, as soon as the user message arrives.
        messages = state.get("messages", [])
        user_messages = [m for m in messages if self._message_type(m) in {"human", "user"}]
        return len(user_messages) == 1

    def _should_generate_title(self, state: TitleMiddlewareState) -> bool:
        """Check if we should generate the final LLM title for this thread."""
        config = self._get_title_config()
        if not config.enabled:
            return False

        # Skip only if we already have a *finalized* title. A provisional
        # placeholder set in before_model should still be upgraded by the LLM.
        if state.get("title") and not state.get("title_provisional"):
            return False

        # Check if this is the first turn (has at least one user message and one assistant response)
        messages = state.get("messages", [])
        if len(messages) < 2:
            return False

        # Count user and assistant messages
        user_messages = [m for m in messages if self._message_type(m) in {"human", "user"}]
        assistant_messages = [m for m in messages if self._message_type(m) in {"ai", "assistant"}]

        # Generate title after first complete exchange
        return len(user_messages) == 1 and len(assistant_messages) >= 1

    def _build_title_prompt(self, state: TitleMiddlewareState) -> tuple[str, str]:
        """Extract user/assistant messages and build the title prompt.

        Returns (prompt_string, user_msg) so callers can use user_msg as fallback.
        """
        config = self._get_title_config()
        messages = state.get("messages", [])

        user_msg_content = next((self._message_content(m) for m in messages if self._message_type(m) in {"human", "user"}), "")
        assistant_msg_content = next((self._message_content(m) for m in messages if self._message_type(m) in {"ai", "assistant"}), "")

        user_msg = self._normalize_content(user_msg_content)
        assistant_msg = self._strip_think_tags(self._normalize_content(assistant_msg_content))

        prompt = config.prompt_template.format(
            max_words=config.max_words,
            user_msg=user_msg[:500],
            assistant_msg=assistant_msg[:500],
        )
        return prompt, user_msg

    def _strip_think_tags(self, text: str) -> str:
        """Remove <think>...</think> blocks emitted by reasoning models (e.g. minimax, GLM, DeepSeek-R1).

        Also handles malformed cases: an unclosed `<think>` (response truncated mid-reasoning)
        or a stray `</think>` (open tag eaten by upstream parsing) — both can otherwise leak
        raw reasoning into the title.
        """
        text = re.sub(r"<think>[\s\S]*?</think>", "", text, flags=re.IGNORECASE)
        text = re.sub(r"<think>[\s\S]*$", "", text, flags=re.IGNORECASE)
        text = re.sub(r"^[\s\S]*?</think>", "", text, flags=re.IGNORECASE)
        return text.strip()

    def _parse_title(self, content: object) -> str:
        """Normalize model output into a clean title string (strips <think> blocks)."""
        config = self._get_title_config()
        title_content = self._strip_think_tags(self._normalize_content(content))
        title = title_content.strip().strip('"').strip("'")
        return title[: config.max_chars] if len(title) > config.max_chars else title

    def _isolated_config(self) -> dict[str, Any]:
        """Build a RunnableConfig that does NOT inherit the parent's streaming callbacks.

        The leak that pushed title-generation reasoning into the live UI stream came
        from inheriting the parent RunnableConfig (and with it the SSE/streaming
        callbacks) for the title model call. Here we deliberately start from an empty
        config with ``callbacks: []`` so the title model's tokens — including any
        ``<think>`` reasoning from a reasoning-only deployment — never reach the
        frontend stream or the run journal as conversation output. The stored title
        is additionally cleaned by ``_parse_title``/``_strip_think_tags``.
        """
        return {
            "run_name": "title_agent",
            "tags": ["middleware:title"],
            "callbacks": [],
        }

    def _fallback_title(self, user_msg: str) -> str:
        config = self._get_title_config()
        fallback_chars = min(config.max_chars, 50)
        if len(user_msg) > fallback_chars:
            return user_msg[:fallback_chars].rstrip() + "..."
        return user_msg if user_msg else "New Conversation"

    def _provisional_title_result(self, state: TitleMiddlewareState) -> dict | None:
        """Set an instant placeholder title from the first user message (no LLM)."""
        if not self._should_set_provisional_title(state):
            return None

        messages = state.get("messages", [])
        user_msg_content = next((self._message_content(m) for m in messages if self._message_type(m) in {"human", "user"}), "")
        user_msg = self._normalize_content(user_msg_content)
        if not user_msg.strip():
            return None
        title = self._fallback_title(user_msg)
        logger.debug("TitleMiddleware: set provisional title %r before first model call", title)
        return {"title": title, "title_provisional": True}

    def _generate_title_result(self, state: TitleMiddlewareState) -> dict | None:
        """Generate a local fallback title without blocking on an LLM call."""
        if not self._should_generate_title(state):
            return None

        _, user_msg = self._build_title_prompt(state)
        return {"title": self._fallback_title(user_msg), "title_provisional": False}

    def _force_disable_thinking(self, model: object) -> None:
        """Force-disable thinking on the title model via `extra_body` overrides.

        Title generation never needs reasoning. The model factory only emits
        disable params when the model config declares `when_thinking_disabled`
        or a matching `when_thinking_enabled` shape — many user-added reasoning
        models (e.g. GLM-4.5/5) only set `supports_thinking: true` and would
        otherwise keep thinking on, which is slow and leaks <think> tags into
        the title prompt response. Injecting both common OpenAI-compatible
        disable shapes covers Zhipu GLM (`thinking.type`) and Qwen/vLLM
        (`chat_template_kwargs.enable_thinking`); other providers ignore
        unknown body fields.
        """
        if not hasattr(model, "extra_body"):
            return
        current = getattr(model, "extra_body", None) or {}
        merged_chat_template_kwargs = {
            **(current.get("chat_template_kwargs") or {}),
            "enable_thinking": False,
        }
        merged = {
            **current,
            "thinking": {**(current.get("thinking") or {}), "type": "disabled"},
            "chat_template_kwargs": merged_chat_template_kwargs,
        }
        try:
            model.extra_body = merged
        except Exception:
            logger.debug("Could not override extra_body to disable thinking on title model", exc_info=True)

    async def _agenerate_title_result(self, state: TitleMiddlewareState) -> dict | None:
        """Polish the title with an LLM, without leaking its reasoning to the UI.

        The title model is invoked with an *isolated* config (``_isolated_config``)
        that does not inherit the parent's streaming callbacks, so even a
        reasoning-only deployment's ``<think>`` tokens never reach the live UI
        stream or the run journal. Thinking is also force-disabled on the model,
        and the final stored title is stripped of any residual ``<think>`` blocks.
        On any failure we fall back to the local question-derived title.
        """
        if not self._should_generate_title(state):
            return None

        config = self._get_title_config()
        prompt, user_msg = self._build_title_prompt(state)

        try:
            model_kwargs: dict[str, Any] = {"thinking_enabled": False, "force_disable_thinking": True}
            if self._app_config is not None:
                model_kwargs["app_config"] = self._app_config
            if config.model_name:
                model = create_chat_model(name=config.model_name, **model_kwargs)
            else:
                model = create_chat_model(**model_kwargs)
            self._force_disable_thinking(model)
            response = await model.ainvoke(prompt, config=self._isolated_config())
            title = self._parse_title(response.content)
            if title:
                return {"title": title, "title_provisional": False}
        except Exception:
            logger.debug("Failed to generate async title via LLM; falling back to local title", exc_info=True)
        return {"title": self._fallback_title(user_msg), "title_provisional": False}

    @override
    def before_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
        return self._provisional_title_result(state)

    @override
    async def abefore_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
        return self._provisional_title_result(state)

    @override
    def after_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
        # Keep the provisional title in place during the main run. The gateway
        # worker finalizes it in a background task after the answer is complete.
        return None

    @override
    async def aafter_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None:
        # Keep the provisional title in place during the main async run so the
        # worker's post-run background finalizer can replace it with the real
        # LLM-generated title without blocking the input box.
        return None


async def finalize_thread_title(
    *,
    checkpointer: Any,
    thread_store: Any | None,
    thread_id: str,
    source_state: dict[str, Any],
    app_config: "AppConfig | None" = None,
) -> bool:
    """Finalize a provisional thread title in the background.

    ``source_state`` must be the checkpoint snapshot captured when the first run
    completed. That avoids races where the user asks a second question before
    the background title task starts, which would otherwise make the current
    checkpoint fail the "first exchange only" predicate and leave the title
    provisional forever.
    """

    log_title_debug("Title finalize: evaluating background title generation for thread %s", thread_id)

    middleware = TitleMiddleware(app_config=app_config)
    source_snapshot = copy.deepcopy(source_state)
    if not middleware._should_generate_title(source_snapshot):
        log_title_debug(
            "Title finalize: skipped for thread %s because source snapshot is not eligible",
            thread_id,
        )
        return False

    log_title_debug("Title finalize: starting background title generation for thread %s", thread_id)
    result = await middleware._agenerate_title_result(source_snapshot)
    if not result:
        log_title_debug("Title finalize: no title result produced for thread %s", thread_id)
        return False

    read_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
    checkpoint_tuple = await checkpointer.aget_tuple(read_config)
    if checkpoint_tuple is None:
        log_title_debug("Title finalize: skipped for thread %s because checkpoint was not found", thread_id)
        return False

    checkpoint: dict[str, Any] = dict(getattr(checkpoint_tuple, "checkpoint", {}) or {})
    metadata: dict[str, Any] = dict(getattr(checkpoint_tuple, "metadata", {}) or {})
    channel_values: dict[str, Any] = dict(checkpoint.get("channel_values", {}) or {})

    # If another path already finalized or manually renamed the title, keep it.
    if channel_values.get("title_provisional") is not True:
        log_title_debug(
            "Title finalize: skipped for thread %s because title is no longer provisional",
            thread_id,
        )
        return False

    channel_values.update(result)
    checkpoint["channel_values"] = channel_values
    metadata["updated_at"] = now_iso()

    write_config = {
        "configurable": {
            "thread_id": thread_id,
            "checkpoint_ns": "",
        }
    }
    new_versions: dict[str, Any] = {}
    get_next_version = getattr(checkpointer, "get_next_version", None)
    if callable(get_next_version):
        versions = dict(checkpoint.get("channel_versions", {}) or {})
        for key in ("title", "title_provisional"):
            next_version = get_next_version(versions.get(key), None)
            versions[key] = next_version
            new_versions[key] = next_version
        checkpoint["channel_versions"] = versions

    await checkpointer.aput(write_config, checkpoint, metadata, new_versions)

    title = result.get("title")
    if title and thread_store is not None:
        await thread_store.update_display_name(thread_id, title, user_id=None)

    log_title_debug("Title finalize: completed for thread %s with title %r", thread_id, title)

    return True