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

import json
import re
import uuid
from typing import Any

from langchain_core.messages import AIMessage

from apis.rest_countries import lookup_country
from ...messages import parse_text_tool_calls
from ...tools import _parse_arguments, truncate
from ..state import CandidateCountry, TodoItem
from .config import MAX_TODOS, RESEARCH_TOOL_NAMES

# Conservative starting shortlist for skilled IT applicants (used when LLM/search fail).
DEFAULT_SKILLED_IT_COUNTRIES: list[CandidateCountry] = [
    {
        "iso2": "CA",
        "name": "Canada",
        "pathway_hint": "Express Entry / Provincial Nominee skilled worker route",
        "label": "Skilled worker - 12-18 mo",
    },
    {
        "iso2": "DE",
        "name": "Germany",
        "pathway_hint": "EU Blue Card / skilled worker residence route",
        "label": "EU Blue Card - 6-12 mo",
    },
    {
        "iso2": "AU",
        "name": "Australia",
        "pathway_hint": "Skilled Independent / State nomination route",
        "label": "Skilled migration - 12-18 mo",
    },
    {
        "iso2": "IE",
        "name": "Ireland",
        "pathway_hint": "Critical Skills Employment Permit route",
        "label": "Critical Skills - 6-12 mo",
    },
]

_COUNTRY_NAME_TO_ISO2 = {
    "canada": "CA",
    "germany": "DE",
    "australia": "AU",
    "ireland": "IE",
    "new zealand": "NZ",
    "united kingdom": "GB",
    "uk": "GB",
    "portugal": "PT",
    "netherlands": "NL",
    "singapore": "SG",
    "united states": "US",
    "usa": "US",
}


def extract_json(text: str) -> dict[str, Any] | None:
    decoder = json.JSONDecoder()
    for start in range(len(text)):
        if text[start] != "{":
            continue
        try:
            parsed, _ = decoder.raw_decode(text[start:])
        except json.JSONDecodeError:
            continue
        if isinstance(parsed, dict):
            return parsed
    return None


def user_text(user_content: str | list[dict[str, Any]]) -> str:
    if isinstance(user_content, str):
        return user_content
    parts = [
        item.get("text", "")
        for item in user_content
        if isinstance(item, dict) and item.get("type") == "text"
    ]
    return "\n".join(part for part in parts if part)


def profile_summary_from_text(profile_text: str) -> str:
    """Build a short profile summary without calling the LLM."""
    lines = [line.strip() for line in profile_text.splitlines() if line.strip()]
    if not lines:
        return truncate(profile_text, 600)
    headline = lines[0]
    bullets = [line.lstrip("- ").strip() for line in lines[1:] if line.strip().startswith("-")]
    if bullets:
        return truncate(f"{headline}. Key constraints: {'; '.join(bullets[:6])}.", 600)
    return truncate(profile_text, 600)


def heuristic_candidate_countries(profile_text: str) -> list[CandidateCountry]:
    """Deterministic shortlist when discovery/planner cannot produce one."""
    text = profile_text.lower()
    candidates = list(DEFAULT_SKILLED_IT_COUNTRIES)

    if any(word in text for word in ("software", "it", "engineer", "developer", "tech")):
        return candidates[:MAX_TODOS]

    if any(word in text for word in ("study", "student", "university")):
        return [
            {
                "iso2": "DE",
                "name": "Germany",
                "pathway_hint": "Student visa / post-study residence route",
                "label": "Study route - 12-24 mo",
            },
            {
                "iso2": "CA",
                "name": "Canada",
                "pathway_hint": "Study permit / PGWP pathway",
                "label": "Study route - 12-24 mo",
            },
            {
                "iso2": "IE",
                "name": "Ireland",
                "pathway_hint": "Study / graduate route",
                "label": "Study route - 12-24 mo",
            },
            {
                "iso2": "AU",
                "name": "Australia",
                "pathway_hint": "Student visa / skilled graduate route",
                "label": "Study route - 12-24 mo",
            },
        ][:MAX_TODOS]

    return candidates[:MAX_TODOS]


def _candidate_from_iso2(iso2: str) -> CandidateCountry | None:
    info = lookup_country(iso2)
    if not info:
        return None
    default = next(
        (item for item in DEFAULT_SKILLED_IT_COUNTRIES if item["iso2"] == iso2.upper()),
        None,
    )
    if default:
        return default
    return {
        "iso2": info["cca2"],
        "name": str(info["name"]),
        "pathway_hint": "Skilled worker / residence pathway",
        "label": "Skilled route - 12-18 mo",
    }


def candidates_from_search_text(text: str) -> list[CandidateCountry]:
    """Extract mentioned countries from search result text."""
    lowered = text.lower()
    found: list[CandidateCountry] = []
    seen: set[str] = set()
    for name, iso2 in _COUNTRY_NAME_TO_ISO2.items():
        if name in lowered and iso2 not in seen:
            candidate = _candidate_from_iso2(iso2)
            if candidate:
                found.append(candidate)
                seen.add(iso2)
    return found


def merge_candidates(
    primary: list[CandidateCountry],
    secondary: list[CandidateCountry],
) -> list[CandidateCountry]:
    merged: list[CandidateCountry] = []
    seen: set[str] = set()
    for item in [*primary, *secondary]:
        iso2 = item["iso2"].upper()
        if iso2 in seen:
            continue
        merged.append(item)
        seen.add(iso2)
        if len(merged) >= MAX_TODOS:
            break
    return merged


def split_todo_label(label: str) -> tuple[str, str]:
    if "—" in label:
        country, methods = label.split("—", 1)
        return country.strip(), methods.strip()
    if " - " in label:
        country, methods = label.split(" - ", 1)
        return country.strip(), methods.strip()
    return label.strip(), "Skilled migration pathway"


def format_todo_label(todo: TodoItem | dict[str, Any]) -> str:
    country = str(todo.get("country") or "").strip()
    methods = str(todo.get("methods") or "").strip()
    if country and methods:
        return f"{country} — {methods}"
    return str(todo.get("title") or country or "Research task").strip()


def todo_research_brief(todo: TodoItem, profile_summary: str) -> str:
    return (
        f"Country: {todo['country']}\n"
        f"Migration methods to research: {todo['methods']}\n\n"
        f"Research the best realistic skilled migration pathway to {todo['country']} "
        f"for this applicant. Focus on {todo['methods']}. Cover eligibility, required "
        f"documents, approximate costs, realistic timeline within 12-18 months, path to "
        f"permanent residence, and risks. Use official government or immigration authority "
        f"sources.\n\nApplicant profile: {profile_summary}"
    )


def normalize_todo(
    raw: dict[str, Any],
    *,
    todo_id: int,
    profile_summary: str,
    candidate: CandidateCountry | None = None,
) -> TodoItem | None:
    country = str(raw.get("country") or "").strip()
    methods = str(raw.get("methods") or "").strip()

    if not country and raw.get("title"):
        country, parsed_methods = split_todo_label(str(raw["title"]))
        methods = methods or parsed_methods

    if candidate:
        country = country or candidate["name"]
        methods = methods or candidate["pathway_hint"]

    if not methods and raw.get("description"):
        methods = truncate(str(raw["description"]).strip(), 220)

    if not country:
        return None
    if not methods:
        methods = "Skilled migration pathway"

    return {"id": todo_id, "country": country, "methods": methods}


def country_todo(
    candidate: CandidateCountry,
    profile_summary: str,
    *,
    todo_id: int,
) -> TodoItem:
    _ = profile_summary
    return {
        "id": todo_id,
        "country": candidate["name"],
        "methods": candidate["pathway_hint"],
    }


def plan_from_candidates(
    candidates: list[CandidateCountry],
    profile_text: str,
    *,
    thinking: str = "",
) -> dict[str, Any]:
    summary = profile_summary_from_text(profile_text)
    todos = [
        country_todo(candidate, summary, todo_id=index + 1)
        for index, candidate in enumerate(candidates[:MAX_TODOS])
    ]
    return {
        "thinking": thinking,
        "countries": [item["iso2"] for item in candidates[:MAX_TODOS]],
        "labels": [item["label"] for item in candidates[:MAX_TODOS]],
        "profile_summary": summary,
        "todos": todos,
    }


def fallback_plan(
    profile_text: str,
    candidates: list[CandidateCountry] | None = None,
) -> dict[str, Any]:
    shortlist = candidates or heuristic_candidate_countries(profile_text)
    return plan_from_candidates(
        shortlist,
        profile_text,
        thinking=(
            "Using a conservative starting shortlist of skilled-worker destinations. "
            "Each country will be researched in parallel."
        ),
    )


def _is_generic_todo(todo: dict[str, Any]) -> bool:
    country = str(todo.get("country") or "").lower()
    methods = str(todo.get("methods") or "").lower()
    title = str(todo.get("title") or "").lower()
    description = str(todo.get("description") or "").lower()
    generic_titles = {"research migration options", "research task"}
    if title in generic_titles or country in generic_titles:
        return True
    if methods == "research realistic migration options for this profile":
        return True
    if "research realistic migration options for this profile" in description:
        return True
    return len(description) > 800 and description.count("\n") >= 4


def normalize_plan(
    plan: dict[str, Any] | None,
    profile_text: str,
    candidates: list[CandidateCountry],
) -> dict[str, Any]:
    """Ensure the plan has 3-4 useful country-specific todos."""
    summary = str((plan or {}).get("profile_summary") or "").strip() or profile_summary_from_text(
        profile_text
    )
    shortlist = candidates[:MAX_TODOS] or heuristic_candidate_countries(profile_text)

    if plan is None:
        return fallback_plan(profile_text, shortlist)

    raw_todos = plan.get("todos") or []
    todos: list[TodoItem] = []
    for index, raw in enumerate(raw_todos[:MAX_TODOS]):
        if not isinstance(raw, dict):
            continue
        if _is_generic_todo(raw):
            continue
        candidate = shortlist[index] if index < len(shortlist) else None
        normalized = normalize_todo(
            raw,
            todo_id=len(todos) + 1,
            profile_summary=summary,
            candidate=candidate,
        )
        if normalized:
            todos.append(normalized)

    if len(todos) < 3:
        todos = [
            country_todo(candidate, summary, todo_id=index + 1)
            for index, candidate in enumerate(shortlist[:MAX_TODOS])
        ]

    countries = [str(code) for code in plan.get("countries") or [] if code]
    labels = [str(label) for label in plan.get("labels") or [] if label]
    if len(countries) != len(todos):
        countries = [item["iso2"] for item in shortlist[: len(todos)]]
        labels = [item["label"] for item in shortlist[: len(todos)]]

    thinking = str(plan.get("thinking") or "").strip()
    if not thinking and len(todos) >= 3:
        thinking = (
            f"Split research into {len(todos)} parallel country tasks based on the "
            f"applicant profile and discovery shortlist."
        )

    return {
        "thinking": thinking,
        "countries": countries,
        "labels": labels,
        "profile_summary": summary,
        "todos": todos,
    }


def _think_tags() -> tuple[str, str]:
    return "<" + "think" + ">", "</" + "think" + ">"


_THINK_OPEN, _THINK_CLOSE = _think_tags()
_THINK_BLOCK_PATTERN = re.compile(
    re.escape(_THINK_OPEN) + r".*?" + re.escape(_THINK_CLOSE),
    re.I | re.DOTALL,
)
_THINK_INNER_PATTERN = re.compile(
    re.escape(_THINK_OPEN) + r"(.*?)" + re.escape(_THINK_CLOSE),
    re.I | re.DOTALL,
)


def _flatten_content_blocks(content: Any) -> str:
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for block in content:
            if isinstance(block, str):
                parts.append(block)
                continue
            if isinstance(block, dict):
                if block.get("type") == "text":
                    parts.append(str(block.get("text") or ""))
                elif "text" in block:
                    parts.append(str(block["text"]))
        return "\n".join(part for part in parts if part)
    return str(content)


def _extract_reasoning(message: AIMessage) -> str:
    reasoning = getattr(message, "reasoning", None)
    if reasoning:
        return str(reasoning).strip()

    additional = getattr(message, "additional_kwargs", None) or {}
    for key in ("reasoning_content", "reasoning", "reasoning_text"):
        value = additional.get(key)
        if value:
            return str(value).strip()

    metadata = getattr(message, "response_metadata", None) or {}
    for key in ("reasoning_content", "reasoning"):
        value = metadata.get(key)
        if value:
            return str(value).strip()
    return ""


def _strip_think_blocks(text: str) -> str:
    if not text.strip():
        return ""
    outside = _THINK_BLOCK_PATTERN.sub("", text).strip()
    if outside:
        return outside
    inner_parts = _THINK_INNER_PATTERN.findall(text)
    if inner_parts:
        return "\n".join(part.strip() for part in inner_parts if part.strip())
    return text.strip()


def assistant_text_sources(message: AIMessage) -> tuple[str, str]:
    content_text = _flatten_content_blocks(message.content).strip()
    reasoning_text = _extract_reasoning(message)
    return content_text, reasoning_text


def extract_assistant_text(message: AIMessage) -> str:
    """Extract user-visible assistant text from content, blocks, or reasoning."""
    content_text, reasoning_text = assistant_text_sources(message)
    for candidate in (content_text, reasoning_text):
        stripped = _strip_think_blocks(candidate)
        if stripped:
            return stripped
    return ""


def extract_thinking_text(message: AIMessage) -> str:
    """Extract model reasoning / thinking text separate from the user-facing answer."""
    content_text, reasoning_text = assistant_text_sources(message)
    answer = extract_assistant_text(message)
    parts: list[str] = []

    for text in (reasoning_text, content_text):
        if not text.strip():
            continue
        inner_parts = _THINK_INNER_PATTERN.findall(text)
        if inner_parts:
            parts.extend(part.strip() for part in inner_parts if part.strip())
            continue
        outside = _THINK_BLOCK_PATTERN.sub("", text).strip()
        candidate = outside or text.strip()
        if candidate and candidate != answer:
            parts.append(candidate)

    deduped: list[str] = []
    seen: set[str] = set()
    for part in parts:
        if part not in seen:
            seen.add(part)
            deduped.append(part)
    return "\n\n".join(deduped).strip()


def research_tool_calls(
    response: AIMessage,
) -> list[tuple[str, dict[str, Any], str]]:
    """Return normalized tool calls, including text-emitted calls from small models."""
    if response.tool_calls:
        return [
            (tool_call["name"], tool_call["args"] or {}, tool_call["id"])
            for tool_call in response.tool_calls
        ]

    content_text, reasoning_text = assistant_text_sources(response)
    parsed_calls = parse_text_tool_calls(content_text) or parse_text_tool_calls(
        reasoning_text
    )
    if not parsed_calls:
        return []

    normalized: list[tuple[str, dict[str, Any], str]] = []
    for tool_call in parsed_calls:
        function = tool_call.get("function") or {}
        tool_name = str(function.get("name") or "")
        if tool_name not in RESEARCH_TOOL_NAMES:
            continue
        normalized.append(
            (
                tool_name,
                _parse_arguments(str(function.get("arguments") or "")),
                str(tool_call.get("id") or uuid.uuid4().hex),
            )
        )
    return normalized


def discovery_queries(profile_text: str) -> list[str]:
    text = re.sub(r"\s+", " ", profile_text).strip()
    occupation = "software engineer" if re.search(r"software|it|developer", text, re.I) else "skilled worker"
    origin = "India" if re.search(r"\bindia\b", text, re.I) else "applicant country"
    return [
        (
            f"best skilled worker immigration pathways {occupation} {origin} "
            "official government permanent residence"
        ),
        (
            f"countries skilled worker visa path to permanent residence "
            f"{occupation} Indian citizen official immigration"
        ),
    ]