# ui/agent/synthesis.py
from __future__ import annotations

import json
import re
from typing import Any, TypedDict


class Finding(TypedDict, total=False):
    todo_id: int
    todo_title: str
    todo_country: str
    todo_methods: str
    summary: str


class TodoItem(TypedDict):
    id: int
    country: str
    methods: str


URL_PATTERN = re.compile(r"https?://[^\s\)\]>\"']+")
STEP_PATTERN = re.compile(
    r"(?:^|\n)\s*(?:\d+[\.\)]|[-*])\s+.+(?:\n\s*(?:\d+[\.\)]|[-*])\s+.+)*",
    re.MULTILINE,
)
SECTION_ALIASES = {
    "eligibility": re.compile(r"(?:^|\n)#+\s*eligibility[^\n]*\n(.+?)(?=\n#+\s|\Z)", re.I | re.S),
    "documents": re.compile(
        r"(?:^|\n)#+\s*(?:required documents|documents)[^\n]*\n(.+?)(?=\n#+\s|\Z)",
        re.I | re.S,
    ),
    "timeline": re.compile(
        r"(?:^|\n)#+\s*(?:timeline|duration|processing)[^\n]*\n(.+?)(?=\n#+\s|\Z)",
        re.I | re.S,
    ),
    "costs": re.compile(
        r"(?:^|\n)#+\s*(?:costs?|fees?|budget)[^\n]*\n(.+?)(?=\n#+\s|\Z)",
        re.I | re.S,
    ),
    "risks": re.compile(
        r"(?:^|\n)#+\s*(?:risks?|tradeoffs?)[^\n]*\n(.+?)(?=\n#+\s|\Z)",
        re.I | re.S,
    ),
    "steps": re.compile(
        r"(?:^|\n)#+\s*(?:steps?|procedure|process)[^\n]*\n(.+?)(?=\n#+\s|\Z)",
        re.I | re.S,
    ),
    "pathway": re.compile(
        r"(?:^|\n)#+\s*(?:visa pathway|pathway|route)[^\n]*\n(.+?)(?=\n#+\s|\Z)",
        re.I | re.S,
    ),
}


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


def _extract_section(text: str, key: str) -> str | None:
    pattern = SECTION_ALIASES.get(key)
    if not pattern:
        return None
    match = pattern.search(text)
    if not match:
        return None
    value = match.group(1).strip()
    return value or None


def _extract_steps(text: str) -> str | None:
    section = _extract_section(text, "steps")
    if section:
        return section
    matches = STEP_PATTERN.findall(text)
    if matches:
        return "\n".join(match.strip() for match in matches[:1])
    return None


def _extract_urls(text: str) -> list[str]:
    seen: set[str] = set()
    urls: list[str] = []
    for url in URL_PATTERN.findall(text):
        cleaned = url.rstrip(".,;")
        if cleaned not in seen:
            seen.add(cleaned)
            urls.append(cleaned)
    return urls


def _needs_verification(value: str | None, *, fallback: str) -> str:
    if value and value.strip():
        return value.strip()
    return f"{fallback} *(Needs verification on official immigration sources.)*"


def build_country_report_sections(
    summary: str,
    *,
    country: str,
    pathway: str,
) -> dict[str, str]:
    """Structured sections for a single country report."""
    pathway_text = _extract_section(summary, "pathway") or pathway
    timeline = _extract_section(summary, "timeline")
    costs = _extract_section(summary, "costs")
    risks = _extract_section(summary, "risks")
    urls = _extract_urls(summary)

    duration_parts = []
    if timeline:
        duration_parts.append(f"**Timeline:** {timeline}")
    if costs:
        duration_parts.append(f"**Costs:** {costs}")
    if risks:
        duration_parts.append(f"**Risks:** {risks}")
    duration_block = "\n\n".join(duration_parts) if duration_parts else ""

    sources_block = "\n".join(f"- {url}" for url in urls[:8]) if urls else ""

    return {
        "why_recommended": _needs_verification(
            _extract_section(summary, "eligibility"),
            fallback=f"Potential fit for skilled migration to {country} via {pathway}.",
        ),
        "feasible_methods": _needs_verification(
            pathway_text,
            fallback=pathway,
        ),
        "procedure": _needs_verification(
            _extract_steps(summary),
            fallback=(
                f"Confirm the official application sequence on the immigration "
                f"authority website for {country}."
            ),
        ),
        "requirements": _needs_verification(
            _extract_section(summary, "documents"),
            fallback=(
                "Passport, education credentials, employment records, language "
                f"test results (if required), and proof of funds — verify exact "
                f"list for {pathway}."
            ),
        ),
        "duration_costs_risks": _needs_verification(
            duration_block or None,
            fallback="Typical processing time, fees, and risks not verified in research.",
        ),
        "official_sources": _needs_verification(
            sources_block or None,
            fallback="Check the country's official immigration authority website.",
        ),
    }


def build_structured_final_answer(
    *,
    profile_summary: str,
    findings: list[Finding],
    todos: list[TodoItem] | None = None,
    preamble: str = "",
) -> str:
    """Build a deterministic immigration report from parallel research findings."""
    ordered_findings = sorted(findings, key=lambda item: item["todo_id"])
    todo_by_id = {todo["id"]: todo for todo in (todos or [])}

    if not ordered_findings and todos:
        ordered_findings = [
            {
                "todo_id": todo["id"],
                "todo_title": f"{todo['country']} — {todo['methods']}",
                "todo_country": todo["country"],
                "todo_methods": todo["methods"],
                "summary": "No detailed findings were captured for this country yet.",
            }
            for todo in sorted(todos, key=lambda item: item["id"])
        ]

    lines: list[str] = [
        "# Migration Research Summary",
        "",
        preamble.strip() or (
            "Based on the research gathered so far, here is a structured comparison "
            "of realistic migration options for your profile."
        ),
        "",
        "## Recommended Countries",
    ]

    country_entries: list[dict[str, str]] = []
    all_urls: list[str] = []

    for finding in ordered_findings:
        todo = todo_by_id.get(finding["todo_id"])
        if todo:
            country = todo["country"]
            pathway = todo["methods"]
            title = f"{country} — {pathway}"
        elif finding.get("todo_country") and finding.get("todo_methods"):
            country = str(finding["todo_country"])
            pathway = str(finding["todo_methods"])
            title = f"{country} — {pathway}"
        else:
            title = str(finding.get("todo_title") or "")
            country, pathway = _split_country_title(title)
        summary = str(finding.get("summary") or "").strip()
        country_entries.append(
            {
                "country": country,
                "pathway": pathway,
                "title": title,
                "summary": summary,
            }
        )
        lines.append(f"- **{country}** — {pathway}")
        all_urls.extend(_extract_urls(summary))

    if not country_entries:
        lines.append("- Needs verification — rerun research to populate country shortlist.")

    lines.extend(["", "## Feasible Migration Pathways"])
    for entry in country_entries:
        pathway_text = _extract_section(entry["summary"], "pathway") or entry["pathway"]
        lines.extend(
            [
                "",
                f"### {entry['country']}",
                _needs_verification(pathway_text, fallback=entry["pathway"]),
            ]
        )
        eligibility = _extract_section(entry["summary"], "eligibility")
        if eligibility:
            lines.extend(["", "**Eligibility notes:**", eligibility])

    lines.extend(["", "## Step-By-Step Procedure"])
    for entry in country_entries:
        steps = _extract_steps(entry["summary"])
        lines.extend(["", f"### {entry['country']}"])
        lines.append(
            _needs_verification(
                steps,
                fallback=(
                    "Confirm the official application sequence on the immigration "
                    f"authority website for {entry['country']}."
                ),
            )
        )

    lines.extend(["", "## Requirements And Documents"])
    for entry in country_entries:
        documents = _extract_section(entry["summary"], "documents")
        lines.extend(["", f"### {entry['country']}"])

        lines.append(
            _needs_verification(
                documents,
                fallback=(
                    "Passport, education credentials, employment records, language "
                    "test results (if required), and proof of funds — verify exact "
                    f"list for {entry['pathway']}."
                ),
            )
        )

    lines.extend(["", "## Duration, Costs, And Risks"])
    for entry in country_entries:
        timeline = _extract_section(entry["summary"], "timeline")
        costs = _extract_section(entry["summary"], "costs")
        risks = _extract_section(entry["summary"], "risks")
        lines.extend(["", f"### {entry['country']}"])
        lines.append(
            "**Timeline:** "
            + _needs_verification(
                timeline,
                fallback="Typical processing time not verified in research.",
            )
        )
        lines.append(
            "**Costs:** "
            + _needs_verification(
                costs,
                fallback="Government fees and proof-of-funds requirements not verified.",
            )
        )
        lines.append(
            "**Risks:** "
            + _needs_verification(
                risks,
                fallback="Policy changes, job-market competition, and credential recognition.",
            )
        )

    lines.extend(["", "## Official Sources To Verify"])
    unique_urls = list(dict.fromkeys(all_urls))
    if unique_urls:
        for url in unique_urls[:20]:
            lines.append(f"- {url}")
    else:
        lines.append(
            "- No official URLs were captured in the research notes. "
            "Check each country's immigration authority website before applying."
        )

    lines.extend(
        [
            "",
            "## Next Steps",
            "1. Pick your top 1-2 countries from the shortlist above.",
            "2. Verify eligibility rules, fees, and document checklists on official sites.",
            "3. Start language tests or credential assessments if required.",
            "4. Gather employment and financial evidence within your budget and timeline.",
            "5. Consult a qualified immigration professional before submitting applications.",
            "",
            "## Profile Snapshot",
            profile_summary or "Profile summary unavailable.",
            "",
            "*This is research guidance, not legal advice. Rules change — verify everything "
            "on official government immigration websites.*",
        ]
    )

    return "\n".join(lines)


def findings_from_ui_messages(ui_messages: list[Any]) -> list[Finding]:
    """Best-effort extraction of country findings from legacy ChatMessage cards."""
    findings: list[Finding] = []
    todo_id = 1
    for message in ui_messages:
        metadata = getattr(message, "metadata", None) or {}
        title = str(metadata.get("title") or "")
        if not title.startswith("Findings ·"):
            continue
        todo_title = title.removeprefix("Findings ·").strip()
        content = str(getattr(message, "content", "") or "").strip()
        findings.append(
            {
                "todo_id": todo_id,
                "todo_title": todo_title or f"Research task {todo_id}",
                "summary": content or "No findings captured.",
            }
        )
        todo_id += 1
    return findings


def _parse_tool_payload(content: str) -> dict[str, Any]:
    try:
        parsed = json.loads(content)
    except json.JSONDecodeError:
        return {"text": content}
    return parsed if isinstance(parsed, dict) else {"text": content}


def synthesize_finding_from_tool_messages(
    todo: TodoItem,
    tool_messages: list[Any],
) -> str:
    """Build a markdown finding from tool results when the LLM summary is empty."""
    country = todo["country"]
    methods = todo["methods"]

    search_snippets: list[str] = []
    scraped_snippets: list[str] = []
    profile_notes: list[str] = []
    urls: list[str] = []

    for message in tool_messages:
        content = getattr(message, "content", message)
        payload = _parse_tool_payload(str(content or ""))
        if payload.get("error") and len(payload) <= 3:
            continue

        for result in payload.get("results") or []:
            if not isinstance(result, dict):
                continue
            url = str(result.get("url") or "")
            if url:
                urls.append(url)
            title = str(result.get("title") or url or "Search result")
            highlights = result.get("highlights") or []
            summary = str(result.get("summary") or "")
            snippet = "\n".join(str(item) for item in highlights[:2])
            if not snippet:
                snippet = summary
            if snippet:
                search_snippets.append(f"- **{title}**: {snippet[:500]}")

        markdown = str(payload.get("markdown") or "")
        if markdown:
            title = str(payload.get("title") or payload.get("url") or "Official page")
            url = str(payload.get("source_url") or payload.get("url") or "")
            if url:
                urls.append(url)
            scraped_snippets.append(f"### {title}\n{markdown[:1500]}")

        for page in payload.get("pages") or []:
            if not isinstance(page, dict):
                continue
            page_markdown = str(page.get("markdown") or "")
            if not page_markdown:
                continue
            title = str(page.get("title") or page.get("url") or "Page")
            url = str(page.get("source_url") or page.get("url") or "")
            if url:
                urls.append(url)
            scraped_snippets.append(f"### {title}\n{page_markdown[:800]}")

        for profile in payload.get("countries") or []:
            if not isinstance(profile, dict):
                continue
            name = str(profile.get("name") or country)
            domains = profile.get("official_immigration_domains") or []
            domain_text = ", ".join(str(domain) for domain in domains[:5])
            if domain_text:
                profile_notes.append(f"{name}: official domains include {domain_text}")

    if not search_snippets and not scraped_snippets and not profile_notes:
        return ""

    raw_parts = [
        f"# Research Findings: {country} — {methods}",
        "",
        "*Compiled from tool results. Verify all details on official sources.*",
        "",
        "## Visa Pathway",
        methods,
        "",
    ]
    if search_snippets:
        raw_parts.extend(["## Search Results", "\n".join(search_snippets[:8]), ""])
    if scraped_snippets:
        raw_parts.extend(
            ["## Official Page Excerpts", "\n\n".join(scraped_snippets[:3]), ""]
        )
    if profile_notes:
        raw_parts.extend(["## Country Context", "\n".join(profile_notes), ""])

    unique_urls = list(dict.fromkeys(url for url in urls if url))
    if unique_urls:
        raw_parts.extend(
            ["## Official Sources", "\n".join(f"- {url}" for url in unique_urls[:10])]
        )

    raw_text = "\n".join(raw_parts)
    sections = build_country_report_sections(
        raw_text,
        country=country,
        pathway=methods,
    )

    return (
        f"# {country} — {methods}\n\n"
        "*Automated fallback report from research tool outputs. "
        "Verify all details on official immigration websites.*\n\n"
        f"## Eligibility\n{sections['why_recommended']}\n\n"
        f"## Visa Pathway\n{sections['feasible_methods']}\n\n"
        f"## Required Documents\n{sections['requirements']}\n\n"
        f"## Steps\n{sections['procedure']}\n\n"
        f"## Timeline, Costs, and Risks\n{sections['duration_costs_risks']}\n\n"
        f"## Official Sources\n{sections['official_sources']}\n\n"
        f"## Raw Research Notes\n{raw_text[:2000]}"
    )