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

import json
from typing import Any

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

from apis.exa import search_immigration
from ..state import AgentState, CandidateCountry
from .helpers import (
    candidates_from_search_text,
    discovery_queries,
    heuristic_candidate_countries,
    merge_candidates,
    profile_summary_from_text,
    user_text,
)


def _format_discovery_summary(
    candidates: list[CandidateCountry],
    search_notes: list[str],
) -> str:
    lines = [
        "Initial destination shortlist for parallel research:",
        "",
    ]
    for item in candidates:
        lines.append(
            f"- {item['name']} ({item['iso2']}): {item['pathway_hint']}"
        )
    if search_notes:
        lines.extend(["", "Discovery search notes:"])
        lines.extend(f"- {note}" for note in search_notes[:4])
    lines.append(
        "\nThese are starting candidates only — each country will be researched "
        "in depth before the final recommendation."
    )
    return "\n".join(lines)


def discovery_node(state: AgentState, config: RunnableConfig) -> dict[str, Any]:
    del config  # discovery uses direct API calls, not the LLM
    writer = get_stream_writer()
    profile_text = user_text(state["user_content"])
    profile_summary = profile_summary_from_text(profile_text)
    heuristic = heuristic_candidate_countries(profile_text)

    search_notes: list[str] = []
    search_blob = ""
    search_log: dict[str, Any] = {"queries": [], "results": []}

    for query in discovery_queries(profile_text):
        search_log["queries"].append(query)
        try:
            result = search_immigration(query=query, num_results=6)
            search_log["results"].append(result)
            titles = [
                str(item.get("title") or "")
                for item in result.get("results", [])
                if isinstance(item, dict)
            ]
            if titles:
                search_notes.append(f"{query} → {titles[0]}")
            search_blob += "\n".join(
                f"{item.get('title', '')} {item.get('url', '')}"
                for item in result.get("results", [])
                if isinstance(item, dict)
            )
        except Exception as exc:
            search_notes.append(f"{query} → search unavailable ({exc})")

    from_search = candidates_from_search_text(search_blob)
    candidates = merge_candidates(from_search, heuristic)
    if not candidates:
        candidates = heuristic

    summary = _format_discovery_summary(candidates, search_notes)
    writer(
        {
            "type": "discovery",
            "summary": summary,
            "candidates": candidates,
            "log": {
                "tool": "destination_discovery",
                "arguments": {"queries": search_log["queries"]},
                "result": truncate_discovery_log(search_log),
            },
        }
    )

    return {
        "profile_summary": profile_summary,
        "candidate_countries": candidates,
        "discovery_summary": summary,
    }


def truncate_discovery_log(search_log: dict[str, Any]) -> Any:
    serialized = json.dumps(search_log, default=str)
    if len(serialized) <= 1500:
        return search_log
    return serialized[:1500] + "\n… (truncated)"