# ui/agent/graph/respond.py
from __future__ import annotations

import os
import time
from typing import Any

import gradio as gr
from gradio import ChatMessage

from ui.globe_commands import apply_update_globe
from ui.globe_details import attach_finding_to_globe, attach_research_progress_to_globe

from ..messages import history_to_api_messages, multimodal_input_to_api_content
from ..config import hub_token_available
from ..respond import (
    AUTH_REQUIRED_MESSAGE,
    SESSION_EXPIRED_MESSAGE,
    _auth_error_message,
)
from ..streaming import yield_response
from ..synthesis import build_structured_final_answer
from .workflow import build_workflow


def _format_plan_summary(todos: list[dict[str, Any]]) -> str:
    count = len(todos)
    if count == 0:
        return "No countries queued for research."
    noun = "country" if count == 1 else "countries"
    return f"{count} {noun} queued for parallel research."


def _research_task_country_methods(
    research_task: dict[str, Any],
) -> tuple[str, str] | None:
    country = str(research_task.get("country") or "").strip()
    methods = str(research_task.get("methods") or "").strip()
    if country and methods:
        return country, methods

    title = str(research_task.get("title") or "").strip()
    if not title:
        return None
    from ui.globe_details import split_country_title

    country, methods = split_country_title(title)
    if country:
        return country, methods or "Skilled migration pathway"
    return None


class _UiState:
    """Maps custom graph events onto the running ChatMessage list."""

    def __init__(self, globe_state: dict[str, Any]):
        self.ui_messages: list[ChatMessage] = []
        self.globe_state = globe_state
        self._pending_by_id: dict[str, int] = {}

    def _attach_research_progress(self, event: dict[str, Any]) -> None:
        research_task = event.get("research_task")
        if not research_task:
            return
        country_methods = _research_task_country_methods(research_task)
        if not country_methods:
            return
        country, methods = country_methods
        log = event.get("log") or {}
        self.globe_state = attach_research_progress_to_globe(
            self.globe_state,
            country=country,
            methods=methods,
            tool_name=str(log.get("tool") or ""),
            args=dict(log.get("arguments") or {}),
            message=str(event.get("message") or ""),
        )

    def handle(self, event: dict[str, Any]) -> bool:
        kind = event.get("type")
        if kind == "thinking":
            thinking_text = str(event.get("text") or "").strip()
            self.ui_messages.append(
                ChatMessage(
                    role="assistant",
                    content=thinking_text,
                    metadata={
                        "title": "Thinking",
                        "status": "done",
                        "display": "thinking",
                        "thinking": thinking_text,
                    },
                )
            )
            return True

        if kind == "discovery":
            self.ui_messages.append(
                ChatMessage(
                    role="assistant",
                    content=str(event.get("summary") or ""),
                    metadata={
                        "title": "Destination discovery",
                        "status": "done",
                        "compact": True,
                        "log": event.get("log") or {},
                    },
                )
            )
            return True

        if kind == "plan":
            todos = event.get("todos") or []
            self.ui_messages.append(
                ChatMessage(
                    role="assistant",
                    content=_format_plan_summary(todos),
                    metadata={
                        "title": "Research plan",
                        "status": "done",
                        "display": "plan",
                        "plan_todos": todos,
                    },
                )
            )
            return True

        if kind == "globe":
            args = event.get("args") or {}
            countries = ", ".join(args.get("countries") or [])
            try:
                result, self.globe_state = apply_update_globe(self.globe_state, args)
                content = (
                    f"Marked {countries or 'recommended countries'} on the globe."
                )
                status = "done"
            except Exception as exc:
                result = {"error": str(exc)}
                content = (
                    f"Could not update globe for {countries or 'recommended countries'}: "
                    f"{exc}"
                )
                status = "done"
            self.ui_messages.append(
                ChatMessage(
                    role="assistant",
                    content=content,
                    metadata={
                        "title": "Updating globe",
                        "status": status,
                        "log": {
                            "tool": "update_globe",
                            "arguments": args,
                            "result": result,
                        },
                    },
                )
            )
            return True

        if kind == "tool_start":
            research_task = event.get("research_task")
            self.ui_messages.append(
                ChatMessage(
                    role="assistant",
                    content=str(event.get("message") or ""),
                    metadata={
                        "title": str(event.get("title") or "Tool"),
                        "status": "pending",
                        "log": event.get("log") or {},
                        "research_task": research_task,
                        "compact": bool(research_task),
                    },
                )
            )
            self._pending_by_id[str(event.get("id"))] = len(self.ui_messages) - 1
            self._attach_research_progress(event)
            return True

        if kind == "tool_end":
            index = self._pending_by_id.pop(str(event.get("id")), None)
            research_task = event.get("research_task")
            message = ChatMessage(
                role="assistant",
                content=str(event.get("message") or ""),
                metadata={
                    "title": str(event.get("title") or "Tool"),
                    "status": "done",
                    "duration": event.get("duration"),
                    "log": event.get("log") or {},
                    "research_task": research_task,
                    "compact": bool(research_task),
                },
            )
            if index is not None and index < len(self.ui_messages):
                self.ui_messages[index] = message
            else:
                self.ui_messages.append(message)
            self._attach_research_progress(event)
            return True

        if kind == "finding":
            research_task = event.get("research_task") or {}
            country_methods = _research_task_country_methods(research_task)
            if country_methods:
                country, methods = country_methods
            else:
                from ui.globe_details import split_country_title

                todo_title = str(event.get("todo_title") or "research")
                country, methods = split_country_title(todo_title)
            todo_label = f"{country} — {methods}"
            summary = str(event.get("summary") or "")
            self.globe_state = attach_finding_to_globe(
                self.globe_state,
                country=country,
                methods=methods,
                summary=summary,
            )
            self.ui_messages.append(
                ChatMessage(
                    role="assistant",
                    content=summary,
                    metadata={
                        "title": f"Findings · {todo_label}",
                        "status": "done",
                        "markdown": True,
                        "research_task": research_task
                        or {
                            "id": event.get("todo_id"),
                            "country": country,
                            "methods": methods,
                        },
                        "research_finding": True,
                    },
                )
            )
            return True

        return False


def respond_with_graph(
    message: str | dict[str, Any],
    history: list[dict[str, Any]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    globe_state: dict[str, Any],
    hf_token: gr.OAuthToken | None,
):
    """LangGraph-backed drop-in replacement for ui.agent.respond.respond."""
    if hf_token is None and not hub_token_available():
        yield from yield_response([], AUTH_REQUIRED_MESSAGE, globe_state)
        return

    if hf_token is not None and hf_token.expires_at <= time.time():
        yield from yield_response([], SESSION_EXPIRED_MESSAGE, globe_state)
        return

    user_content = multimodal_input_to_api_content(message)
    if not user_content:
        yield from yield_response(
            [],
            "Please enter a message or attach a file.",
            globe_state,
        )
        return

    workflow = build_workflow()
    ui_state = _UiState(globe_state)
    final_answer = ""
    last_state: dict[str, Any] = {}

    config = {
        "configurable": {
            "hf_token": (
                hf_token.token
                if hf_token is not None
                else os.environ.get("HF_TOKEN", "")
            ),
            "max_tokens": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
        },
        "recursion_limit": 50,
    }
    graph_input = {
        "user_content": user_content,
        "history_messages": history_to_api_messages(history),
        "findings": [],
    }

    try:
        for mode, payload in workflow.stream(
            graph_input,
            config=config,
            stream_mode=["custom", "values"],
        ):
            if mode == "custom":
                if ui_state.handle(payload):
                    yield list(ui_state.ui_messages), ui_state.globe_state
            elif mode == "values":
                final_answer = payload.get("final_answer") or final_answer
                last_state = payload
    except Exception as exc:
        error_message = _auth_error_message(exc) or (
            f"Sorry, something went wrong while generating a response: {exc}"
        )
        yield from yield_response(
            ui_state.ui_messages, error_message, ui_state.globe_state
        )
        return

    if not final_answer:
        final_answer = build_structured_final_answer(
            profile_summary=str(last_state.get("profile_summary") or "").strip(),
            findings=last_state.get("findings") or [],
            todos=last_state.get("todos"),
            preamble=(
                "I could not complete model synthesis, but here is the structured "
                "research gathered so far."
            ),
        )

    yield from yield_response(
        ui_state.ui_messages,
        final_answer,
        ui_state.globe_state,
        assistant_metadata={
            "markdown": True,
            "thinking": str(last_state.get("consolidator_thinking") or "").strip(),
        },
    )