"""Load curated System Run case studies for the Streamlit demo.

Labels and section structure follow the SciPaths Dev · CodeAgent Case Studies HTML report.
"""

from __future__ import annotations

import json
from functools import lru_cache
from pathlib import Path
from typing import Any, Optional

CASES_DIR = Path(__file__).resolve().parent / "system_run_cases"
ANALYSIS_PATH = CASES_DIR / "case_study_analysis.json"

SYSTEM_RUN_METHODS: list[dict[str, str]] = [
    {
        "id": "codeagent_parametric",
        "label": "Parametric - LLM Only",
        "banner_title": "1 · codeagent_parametric",
        "banner_desc": "Parametric-only generation (no search/retrieval). Mean F1 ≈ 0.20.",
    },
    {
        "id": "codeagent_websearch_deep",
        # Audience label: "Crawl" = deep_web_search_tool (page/PDF fetch; no separate crawl tool).
        "label": "Websearch Deep - LLM + Web Search + Crawl",
        "banner_title": "2 · codeagent_websearch_deep",
        "banner_desc": (
            "Tools: web_search_tool + deep_web_search_tool (page fetch / crawl). "
            "Mean F1 ≈ 0.32."
        ),
    },
]

DEFAULT_SYSTEM_METHOD = SYSTEM_RUN_METHODS[0]["id"]
DEFAULT_SYSTEM_CASE = "success"

# Case keys exposed in the UI — labels match the HTML report.
_METHOD_CASES: dict[str, list[tuple[str, str]]] = {
    "codeagent_parametric": [("success", "Success"), ("failure", "Failure")],
    "codeagent_websearch_deep": [("success", "Success"), ("failure", "Failure")],
}

CASE_OUTCOME_LABEL: dict[str, str] = {
    "success": "Success",
    "failure": "Failure",
}

TOOLING_NOTE_HTML = (
    "<strong>Tooling by setting.</strong>"
    "<br/>· <code>parametric</code>: parametric knowledge only (no search/retrieval)."
    "<br/>· <code>websearch_deep</code>: <code>web_search_tool</code> + "
    "<code>deep_web_search_tool</code> (URL page fetch)."
)


@lru_cache(maxsize=1)
def load_analysis() -> dict[str, Any]:
    with ANALYSIS_PATH.open("r", encoding="utf-8") as f:
        data = json.load(f)
    if not isinstance(data, dict):
        raise ValueError(f"Invalid case study analysis at {ANALYSIS_PATH}")
    return data


def method_ids() -> list[str]:
    return [m["id"] for m in SYSTEM_RUN_METHODS]


def method_meta(method_id: str) -> dict[str, str]:
    for m in SYSTEM_RUN_METHODS:
        if m["id"] == method_id:
            return m
    return {"id": method_id, "label": method_id, "banner_title": method_id, "banner_desc": ""}


def method_label(method_id: str) -> str:
    return method_meta(method_id).get("label") or method_id


def case_chip_label(method_id: str, case_key: str) -> str:
    """Top-of-card label, e.g. 'Parametric · Success'."""
    pretty = {
        "codeagent_parametric": "Parametric",
        "codeagent_websearch_deep": "WebSearch Deep",
    }.get(method_id, method_label(method_id))
    outcome = CASE_OUTCOME_LABEL.get(case_key, case_key)
    return f"{pretty} · {outcome}"


def list_cases(method_id: str) -> list[tuple[str, str]]:
    return list(_METHOD_CASES.get(method_id, []))


def default_case_for_method(method_id: str) -> str:
    cases = list_cases(method_id)
    return cases[0][0] if cases else DEFAULT_SYSTEM_CASE


def outcome_tone(case_key: str, f1: float) -> str:
    if case_key == "success":
        return "ok"
    if case_key == "failure":
        return "bad"
    if f1 >= 0.75:
        return "ok"
    if f1 <= 0.05:
        return "bad"
    return "part"


def _setting(method_id: str) -> dict[str, Any]:
    settings = load_analysis().get("settings") or {}
    block = settings.get(method_id)
    if not isinstance(block, dict):
        raise KeyError(f"Unknown system-run method: {method_id}")
    return block


def _raw_case_block(method_id: str, case_key: str) -> dict[str, Any]:
    setting = _setting(method_id)
    if case_key == "featured":
        block = setting.get("case")
    else:
        block = setting.get(case_key)
    if not isinstance(block, dict):
        raise KeyError(f"Unknown case {case_key!r} for method {method_id}")
    return block


def _gold_ingredients(case: dict[str, Any]) -> list[dict[str, Any]]:
    """Gold column with full/partial/miss badges and match notes (HTML report style)."""
    predicted = case.get("predicted_ingredients") or []
    gold: list[dict[str, Any]] = []
    for judgment in case.get("recall_judgments") or []:
        if not isinstance(judgment, dict):
            continue
        desc = str(judgment.get("reference_ingredient") or "").strip()
        role = str(judgment.get("reference_role") or "").strip()
        if not desc:
            continue
        covered = bool(judgment.get("covered"))
        partials = judgment.get("partial_matches") or []
        if not isinstance(partials, list):
            partials = []
        match_notes: list[str] = []
        judge_note = ""
        pair_pred_idx: Optional[int] = None
        if covered:
            badge = "full"
            best = judgment.get("best_match") if isinstance(judgment.get("best_match"), dict) else {}
            best_desc = str(best.get("description") or "").strip()
            if best_desc:
                match_notes.append(f"Matched to pred: {best_desc}")
            judge_note = str(judgment.get("reasoning") or "").strip()
            raw_idx = judgment.get("best_match_idx")
            if isinstance(raw_idx, int) and raw_idx > 0:
                pair_pred_idx = raw_idx
            elif best_desc:
                # Fall back to description match against predicted list (1-based).
                for i, pred in enumerate(predicted, start=1):
                    if isinstance(pred, dict) and str(pred.get("description") or "").strip() == best_desc:
                        pair_pred_idx = i
                        break
        elif partials:
            badge = "partial"
            for pm in partials:
                if not isinstance(pm, dict):
                    continue
                pred_idx = pm.get("predicted_idx")
                reason = str(pm.get("reasoning") or "").strip()
                if pred_idx is not None and reason:
                    match_notes.append(f"Partial ↔ pred #{pred_idx}: {reason}")
                elif reason:
                    match_notes.append(reason)
        else:
            badge = "miss"
        gold.append(
            {
                "description": desc,
                "role": role,
                "badge": badge,
                "match_notes": match_notes,
                "judge_note": judge_note,
                "pair_pred_idx": pair_pred_idx,
            }
        )
    return gold


def build_pair_rows(
    gold: list[dict[str, Any]], predicted: list[dict[str, Any]]
) -> list[dict[str, Any]]:
    """Align gold↔prediction rows for success-case connector lines."""
    by_idx = {
        int(p["idx"]): p
        for p in predicted
        if isinstance(p.get("idx"), int)
    }
    used: set[int] = set()
    rows: list[dict[str, Any]] = []
    for g in gold:
        pred_idx = g.get("pair_pred_idx")
        pred = by_idx.get(int(pred_idx)) if isinstance(pred_idx, int) else None
        linked = bool(pred is not None and g.get("badge") == "full")
        if linked and isinstance(pred_idx, int):
            used.add(pred_idx)
        rows.append({"gold": g, "pred": pred, "linked": linked})
    for p in predicted:
        idx = p.get("idx")
        if isinstance(idx, int) and idx not in used:
            rows.append({"gold": None, "pred": p, "linked": False})
    return rows


def _tool_panel(method_id: str, case_key: str, case: dict[str, Any], setting: dict[str, Any]) -> dict[str, Any]:
    """Structured tool evidence / traces for the case card."""
    if method_id == "codeagent_retrieval_plus_websearch":
        analysis = case.get("gemini_analysis") if isinstance(case.get("gemini_analysis"), dict) else {}
        return {
            "kind": "retrieval",
            "intro": (
                "This run used both retrieve_documents and web_search_tool. "
                "Web search continued in later steps after documents were returned."
            ),
            "queries": list(analysis.get("web_search_queries_example") or []),
            "retrieved_docs": list(case.get("retrieved_doc_titles") or []),
        }

    if method_id == "codeagent_websearch_deep":
        traces = setting.get("tool_traces") if isinstance(setting.get("tool_traces"), dict) else {}
        trace = traces.get("success" if case_key == "success" else "failure")
        if not isinstance(trace, dict):
            return {"kind": "none"}
        deep_fetches = []
        for item in trace.get("deep_fetches") or []:
            if not isinstance(item, dict):
                continue
            outcome = str(item.get("outcome") or "")
            ok = outcome.upper().startswith("SUCCESS")
            deep_fetches.append(
                {
                    "url": str(item.get("url") or ""),
                    "outcome": outcome,
                    "ok": ok,
                }
            )
        return {
            "kind": "deep",
            "queries": list(
                trace.get("web_search_queries")
                or trace.get("web_search_queries_sample")
                or []
            ),
            "hits": list(
                trace.get("web_search_hits_notable")
                or trace.get("web_search_correct_hits_but_underused")
                or []
            ),
            "deep_fetches": deep_fetches,
        }

    return {"kind": "none"}


def get_case(method_id: str, case_key: str) -> dict[str, Any]:
    """Normalized case card payload for the UI."""
    setting = _setting(method_id)
    case = _raw_case_block(method_id, case_key)
    tools = [str(t) for t in (setting.get("tools") or [])]
    predicted = []
    for idx, item in enumerate(case.get("predicted_ingredients") or [], start=1):
        if not isinstance(item, dict):
            continue
        predicted.append(
            {
                "idx": idx,
                "description": str(item.get("description") or "").strip(),
                "role": str(item.get("role") or "").strip(),
                "rationale": str(item.get("rationale") or "").strip(),
            }
        )
    f1 = float(case.get("f1") or 0.0)
    return {
        "method_id": method_id,
        "method_label": method_label(method_id),
        "case_key": case_key,
        "case_chip": case_chip_label(method_id, case_key),
        "outcome_tone": outcome_tone(case_key, f1),
        "tools": tools,
        "paper_id": str(case.get("paper_id") or ""),
        "claim_idx": case.get("claim_idx"),
        "paper_title": str(case.get("paper_title") or ""),
        "claim": str(case.get("claim") or ""),
        "recall": float(case.get("recall") or 0.0),
        "precision": float(case.get("precision") or 0.0),
        "f1": f1,
        "gold_ingredients": _gold_ingredients(case),
        "predicted_ingredients": predicted,
        "tool_panel": _tool_panel(method_id, case_key, case, setting),
        "banner": method_meta(method_id),
    }


def get_analysis(method_id: str, case_key: Optional[str] = None) -> dict[str, Any]:
    setting = _setting(method_id)
    if method_id == "codeagent_retrieval_plus_websearch":
        case = setting.get("case") or {}
        analysis = case.get("gemini_analysis") if isinstance(case, dict) else None
        return dict(analysis) if isinstance(analysis, dict) else {}

    analysis = setting.get("gemini_analysis")
    if not isinstance(analysis, dict):
        return {}
    out = dict(analysis)
    out["_case_key"] = case_key or ""
    return out


def analysis_view_for_case(method_id: str, case_key: str) -> dict[str, Any]:
    """Structured analysis view matching HTML report headings."""
    analysis = get_analysis(method_id, case_key)
    if not analysis:
        return {"model_tag": "Gemini 3.1 Pro analysis", "blocks": []}

    blocks: list[dict[str, Any]] = []

    if method_id == "codeagent_retrieval_plus_websearch":
        headline = str(analysis.get("headline") or "").strip()
        takeaway = str(analysis.get("takeaway") or "").strip()
        diag = headline
        if takeaway:
            diag = f"{headline}\n{takeaway}" if headline else takeaway
        if diag:
            blocks.append({"type": "diagnosis", "text": diag, "strong_first": True})
        two_col = []
        for key, title in (
            ("what_was_retrieved", "What the tools returned"),
            ("effect_on_prediction", "Effect on prediction"),
        ):
            items = analysis.get(key)
            if isinstance(items, list) and items:
                two_col.append({"title": title, "items": [str(x) for x in items]})
        if two_col:
            blocks.append({"type": "two_col", "cols": two_col})
        errs = analysis.get("error_taxonomy_bullets")
        if isinstance(errs, list) and errs:
            blocks.append(
                {"type": "list", "title": "Error analysis", "items": [str(x) for x in errs]}
            )
        return {"model_tag": "Gemini 3.1 Pro analysis", "blocks": blocks}

    if method_id == "codeagent_parametric":
        if case_key == "success":
            why = analysis.get("success_why")
            if isinstance(why, list) and why:
                blocks.append(
                    {
                        "type": "list",
                        "title": "Why this succeeded",
                        "items": [str(x) for x in why],
                    }
                )
        else:
            if analysis.get("failure_diagnosis"):
                blocks.append(
                    {
                        "type": "diagnosis",
                        "text": f"Diagnosis. {analysis['failure_diagnosis']}",
                        "strong_first": True,
                    }
                )
            errs = analysis.get("failure_error_bullets")
            if isinstance(errs, list) and errs:
                blocks.append(
                    {
                        "type": "list",
                        "title": "Error taxonomy",
                        "items": [str(x) for x in errs],
                    }
                )
            needed = analysis.get("failure_what_needed")
            if isinstance(needed, list) and needed:
                blocks.append(
                    {
                        "type": "list",
                        "title": "What recovery would have required",
                        "items": [str(x) for x in needed],
                    }
                )
        return {"model_tag": "Gemini 3.1 Pro analysis", "blocks": blocks}

    # websearch_deep
    if case_key == "success":
        two_col = []
        for key, title in (
            ("success_search_contribution", "web_search_tool"),
            ("success_deep_contribution", "deep_web_search_tool"),
        ):
            items = analysis.get(key)
            if isinstance(items, list) and items:
                two_col.append({"title": title, "items": [str(x) for x in items]})
        if two_col:
            blocks.append({"type": "two_col", "cols": two_col})
        why = analysis.get("success_why_overall")
        if isinstance(why, list) and why:
            blocks.append(
                {"type": "list", "title": "Why this succeeded", "items": [str(x) for x in why]}
            )
    else:
        if analysis.get("failure_diagnosis"):
            blocks.append(
                {
                    "type": "diagnosis",
                    "text": f"Diagnosis. {analysis['failure_diagnosis']}",
                    "strong_first": True,
                }
            )
        two_col = []
        for key, title in (
            ("failure_search_what_went_wrong", "web_search_tool"),
            ("failure_deep_what_went_wrong", "deep_web_search_tool"),
        ):
            items = analysis.get(key)
            if isinstance(items, list) and items:
                two_col.append({"title": title, "items": [str(x) for x in items]})
        if two_col:
            blocks.append({"type": "two_col", "cols": two_col})
        errs = analysis.get("failure_error_bullets")
        if isinstance(errs, list) and errs:
            blocks.append(
                {"type": "list", "title": "Error taxonomy", "items": [str(x) for x in errs]}
            )
        needed = analysis.get("failure_what_needed")
        if isinstance(needed, list) and needed:
            blocks.append(
                {
                    "type": "list",
                    "title": "What recovery would have required",
                    "items": [str(x) for x in needed],
                }
            )

    return {"model_tag": "Gemini 3.1 Pro analysis", "blocks": blocks}