# ui/agent/graph/nodes/planner.py
from __future__ import annotations

from typing import Any

from langchain_core.runnables import RunnableConfig
from langgraph.config import get_stream_writer
from langgraph.types import Send

from ..llm import build_llm
from ..state import AgentState, CandidateCountry, TodoItem
from .config import MAX_TODOS, PLANNER_MAX_TOKENS, PLANNER_TEMPERATURE
from .helpers import (
    extract_assistant_text,
    extract_json,
    heuristic_candidate_countries,
    normalize_plan,
    user_text,
)
from .prompts import PLANNER_SYSTEM_PROMPT


def _format_candidate_shortlist(candidates: list[CandidateCountry]) -> str:
    lines = ["Discovery shortlist (preferred starting countries):"]
    for item in candidates:
        lines.append(
            f"- {item['iso2']} {item['name']}: {item['pathway_hint']} ({item['label']})"
        )
    return "\n".join(lines)


def planner_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
    writer = get_stream_writer()
    llm = build_llm(
        config,
        max_tokens=PLANNER_MAX_TOKENS,
        temperature=PLANNER_TEMPERATURE,
    )

    profile_text = user_text(state["user_content"])
    candidates = state.get("candidate_countries") or heuristic_candidate_countries(
        profile_text
    )
    discovery_summary = str(state.get("discovery_summary") or "").strip()
    profile_summary = str(state.get("profile_summary") or "").strip()

    planner_context = "\n\n".join(
        part
        for part in [
            _format_candidate_shortlist(candidates),
            f"Discovery notes:\n{discovery_summary}" if discovery_summary else "",
            f"Profile summary:\n{profile_summary}" if profile_summary else "",
        ]
        if part
    )

    messages: list[Any] = [
        {"role": "system", "content": PLANNER_SYSTEM_PROMPT},
        *state.get("history_messages", []),
        {
            "role": "user",
            "content": (
                f"{planner_context}\n\n"
                f"Original user request:\n{profile_text}\n\n"
                "Produce the JSON research plan now."
            ),
        },
    ]

    raw_plan: dict[str, Any] | None = None
    for _ in range(2):
        response = llm.invoke(messages)
        raw_plan = extract_json(extract_assistant_text(response))
        if raw_plan and raw_plan.get("todos"):
            break
        raw_plan = None

    plan = normalize_plan(raw_plan, profile_text, candidates)

    thinking = str(plan.get("thinking") or "").strip()
    if thinking:
        writer({"type": "thinking", "text": thinking})

    todos: list[TodoItem] = list(plan.get("todos") or [])
    if not todos:
        plan = normalize_plan(None, profile_text, candidates)
        todos = plan["todos"]

    writer({"type": "plan", "todos": todos})

    countries = [str(code) for code in plan.get("countries") or [] if code]
    if countries:
        labels = [str(label) for label in plan.get("labels") or []]
        writer(
            {
                "type": "globe",
                "args": {"action": "show", "countries": countries, "labels": labels},
            }
        )

    return {
        "todos": todos,
        "profile_summary": str(plan.get("profile_summary") or profile_summary),
    }


def fan_out_research(state: AgentState) -> list[Send]:
    return [
        Send(
            "researcher",
            {"todo": todo, "profile_summary": state.get("profile_summary", "")},
        )
        for todo in state["todos"]
    ]