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

from typing import Any

from langchain_core.runnables import RunnableConfig

from ...tools import truncate
from ..llm import build_llm
from ..state import AgentState
from ...synthesis import build_structured_final_answer
from .config import CONSOLIDATOR_MAX_TOKENS, FINDING_SUMMARY_LIMIT
from .helpers import extract_assistant_text, extract_thinking_text, format_todo_label
from .prompts import CONSOLIDATOR_SYSTEM_PROMPT


def consolidator_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
    findings = sorted(state.get("findings", []), key=lambda item: item["todo_id"])
    todo_by_id = {todo["id"]: todo for todo in state.get("todos") or []}
    findings_text = "\n\n".join(
        f"### Finding {item['todo_id']}: "
        f"{format_todo_label(todo_by_id[item['todo_id']]) if item['todo_id'] in todo_by_id else item.get('todo_title', 'Research')}\n"
        f"{truncate(item['summary'], FINDING_SUMMARY_LIMIT)}"
        for item in findings
    )

    llm = build_llm(config, max_tokens=CONSOLIDATOR_MAX_TOKENS)
    messages: list[Any] = [
        {"role": "system", "content": CONSOLIDATOR_SYSTEM_PROMPT},
        *state.get("history_messages", []),
        {"role": "user", "content": state["user_content"]},
        {
            "role": "user",
            "content": (
                "Research team findings:\n\n"
                f"{findings_text or 'No findings were produced.'}\n\n"
                "Consolidate these findings into the final answer now, following "
                "the required format. If findings are incomplete, still produce "
                "the full structured answer and mark missing details as needing "
                "verification."
            ),
        },
    ]
    response = llm.invoke(messages)
    answer = extract_assistant_text(response)
    thinking = extract_thinking_text(response)
    if not answer:
        answer = build_structured_final_answer(
            profile_summary=str(state.get("profile_summary") or "").strip(),
            findings=findings,
            todos=state.get("todos"),
            preamble=(
                "Model synthesis was unavailable, so this answer was assembled "
                "directly from the parallel country research notes below."
            ),
        )
    return {"final_answer": answer, "consolidator_thinking": thinking}