"""Gradio app for the Amul AI / CeRAI evaluation submission.

Tabs:
  1. Report             — the written findings, rendered.
  2. Live demo          — try Amul AI's chat endpoint right here.
  3. CeRAI run results  — table + per-dataset breakdown from the latest CeRAI run.
  4. Run mini-eval      — kick off a small eval against Amul AI in-browser; uses
                          a built-in 6-prompt slice from the BBK dairy split so
                          the Space stays fast.

Designed to work both locally (`python app.py`) and on Hugging Face Spaces.
On Spaces, set the model `Hardware` to CPU basic — no GPU is needed.
"""
from __future__ import annotations

import json
import os
import re
import time
import uuid
from pathlib import Path
from threading import Lock

import gradio as gr
import httpx
import plotly.graph_objects as go

ROOT = Path(__file__).resolve().parent
REPORT_DIR = ROOT / "report"
RESULTS_DIR = ROOT / "results"

AMUL_API_BASE = "https://api.prod.amulai.in"
AMUL_ORIGIN = "https://amulai.in"
AMUL_USER_AGENT = (
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/120 Safari/537.36"
)


# --------------------------------------------------------------------------- #
# Anonymous JWT cache. Same idea as amul_proxy/server.py but inline so the
# Space doesn't need a sidecar process.
# --------------------------------------------------------------------------- #
class _TokenCache:
    def __init__(self) -> None:
        self._token: str | None = None
        self._expiry: float = 0.0
        self._lock = Lock()

    def get(self) -> str:
        with self._lock:
            if self._token and time.time() < self._expiry - 60:
                return self._token
            r = httpx.post(
                f"{AMUL_API_BASE}/api/auth/anonymous",
                headers={
                    "origin": AMUL_ORIGIN,
                    "referer": f"{AMUL_ORIGIN}/",
                    "content-type": "application/json",
                    "user-agent": AMUL_USER_AGENT,
                },
                json={},
                timeout=30,
            )
            r.raise_for_status()
            data = r.json()
            self._token = data["access_token"]
            self._expiry = time.time() + int(data.get("expires_in", 3600))
            return self._token


TOKEN = _TokenCache()


def call_amul(query: str, session_id: str | None = None,
              source_lang: str = "en", target_lang: str = "en") -> tuple[str, float]:
    """One round-trip to Amul AI. Returns (text, elapsed_seconds)."""
    sid = session_id or f"hfspace-{uuid.uuid4()}"
    started = time.time()
    token = TOKEN.get()
    params = {
        "session_id": sid,
        "query": query,
        "source_lang": source_lang,
        "target_lang": target_lang,
        "use_translation_pipeline": "false",
    }
    headers = {
        "authorization": f"Bearer {token}",
        "accept": "*/*",
        "origin": AMUL_ORIGIN,
        "referer": f"{AMUL_ORIGIN}/",
        "user-agent": AMUL_USER_AGENT,
    }
    chunks: list[str] = []
    with httpx.stream("GET", f"{AMUL_API_BASE}/api/chat/", params=params,
                      headers=headers, timeout=120) as r:
        r.raise_for_status()
        for piece in r.iter_text():
            if piece:
                chunks.append(piece)
    return "".join(chunks).strip(), time.time() - started


# --------------------------------------------------------------------------- #
# Tab content                                                                  #
# --------------------------------------------------------------------------- #
def _read_or_placeholder(path: Path, fallback: str) -> str:
    return path.read_text() if path.exists() else fallback


def _latest_run_dir() -> Path | None:
    if not RESULTS_DIR.exists():
        return None
    runs = sorted(
        (p for p in RESULTS_DIR.iterdir() if (p / "summary.json").exists()),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    return runs[0] if runs else None


# Each of the four written deliverables under report/ becomes its own tab.
# Re-read at build_ui() time, so editing the markdown and re-launching the
# Space (or running `python app.py`) picks up changes without a code edit.
REPORT_DOCS: list[tuple[str, str, str]] = [
    # (tab label,            file basename,                                  fallback)
    ("Findings",            "findings.md",                       "_missing report/findings.md_"),
    ("How CeRAI works",     "cerai_overview.md",                 "_missing report/cerai_overview.md_"),
    ("Edits we made",       "cerai_ai_evaluation_tool_edits.md", "_missing report/cerai_ai_evaluation_tool_edits.md_"),
    ("Proposed extensions", "extension_proposal.md",             "_missing report/extension_proposal.md_"),
]


# --------------------------------------------------------------------------- #
# Charts                                                                       #
# --------------------------------------------------------------------------- #
# Three plots that carry the headline story from results/<latest>/*.json:
#   1. Surface metrics — BBK vs KCC bars, colour-coded by metric.
#   2. Surface metrics vs LLM judge — the "judge sees what BLEU misses" story.
#   3. LLM judge score histogram — the bimodal 0/1 distribution.
# Each one degrades to a placeholder if the underlying JSON is missing, so the
# Space still renders cleanly on a fresh clone without a CeRAI run.

_BBK_COLOR = "#4c72b0"  # blue
_KCC_COLOR = "#dd8452"  # orange
_JUDGE_COLOR = "#55a868"  # green
_BG = "rgba(0,0,0,0)"


def _load_summary_and_judge() -> tuple[dict | None, dict | None]:
    rd = _latest_run_dir()
    if not rd:
        return None, None
    summary = json.loads((rd / "summary.json").read_text())
    judge_path = rd / "llm_judge.json"
    judge = json.loads(judge_path.read_text()) if judge_path.exists() else None
    return summary, judge


def _placeholder_figure(text: str) -> go.Figure:
    fig = go.Figure()
    fig.add_annotation(text=text, xref="paper", yref="paper", x=0.5, y=0.5,
                       showarrow=False, font={"size": 14, "color": "#888"})
    fig.update_layout(paper_bgcolor=_BG, plot_bgcolor=_BG,
                      xaxis={"visible": False}, yaxis={"visible": False},
                      margin={"l": 20, "r": 20, "t": 30, "b": 20})
    return fig


def chart_surface_metrics_per_dataset() -> go.Figure:
    summary, _ = _load_summary_and_judge()
    if not summary or not summary.get("by_dataset_metric"):
        return _placeholder_figure("No run results yet — run bash scripts/run_full_eval.sh")

    # Keep only the four content-quality metrics (drop TAT/Error_Rate which
    # aren't on a 0-1 scale and don't tell the surface-overlap story).
    metric_order = ["Lexical_Diversity", "BLEU", "ROUGE", "METEOR"]
    bbk = summary["by_dataset_metric"].get("bbk", {})
    kcc = summary["by_dataset_metric"].get("kcc", {})

    fig = go.Figure(data=[
        go.Bar(name="BBK (structured benchmark)",
               x=metric_order,
               y=[bbk.get(m, {}).get("score_mean", 0) or 0 for m in metric_order],
               marker_color=_BBK_COLOR,
               text=[f"{bbk.get(m, {}).get('score_mean', 0):.2f}" for m in metric_order],
               textposition="outside"),
        go.Bar(name="KCC (real farmer Q&A)",
               x=metric_order,
               y=[kcc.get(m, {}).get("score_mean", 0) or 0 for m in metric_order],
               marker_color=_KCC_COLOR,
               text=[f"{kcc.get(m, {}).get('score_mean', 0):.2f}" for m in metric_order],
               textposition="outside"),
    ])
    fig.update_layout(
        title="CeRAI surface-metric scores: BBK vs KCC",
        barmode="group", yaxis={"range": [0, 1.05], "title": "Mean score (0–1, higher = better)"},
        paper_bgcolor=_BG, plot_bgcolor=_BG,
        legend={"orientation": "h", "yanchor": "bottom", "y": 1.02, "xanchor": "right", "x": 1},
        margin={"l": 50, "r": 20, "t": 80, "b": 50}, height=380,
    )
    return fig


def chart_surface_vs_judge() -> go.Figure:
    summary, judge = _load_summary_and_judge()
    if not summary or not judge:
        return _placeholder_figure("Run the eval + Gemini judge first (steps 5 & 6 in the README)")

    bbk = summary["by_dataset_metric"].get("bbk", {})
    kcc = summary["by_dataset_metric"].get("kcc", {})

    # Mean of BLEU/ROUGE/METEOR per dataset — collapsed into one "surface metric"
    # bar so the comparison against the judge is one-vs-one per dataset.
    def _surface_mean(dat: dict) -> float:
        xs = [dat.get(m, {}).get("score_mean", 0) or 0 for m in ("BLEU", "ROUGE", "METEOR")]
        return sum(xs) / max(1, len(xs))

    surface = {"BBK": _surface_mean(bbk), "KCC": _surface_mean(kcc)}
    judge_means = judge.get("by_dataset", {})
    judge_vals = {"BBK": judge_means.get("bbk", {}).get("mean", 0) or 0,
                  "KCC": judge_means.get("kcc", {}).get("mean", 0) or 0}

    datasets = ["BBK", "KCC"]
    fig = go.Figure(data=[
        go.Bar(name="CeRAI surface metrics (mean of BLEU/ROUGE/METEOR)",
               x=datasets, y=[surface[d] for d in datasets],
               marker_color=_BBK_COLOR,
               text=[f"{surface[d]:.3f}" for d in datasets], textposition="outside"),
        go.Bar(name=f"LLM-as-judge ({judge.get('judge_model', 'gemini')})",
               x=datasets, y=[judge_vals[d] for d in datasets],
               marker_color=_JUDGE_COLOR,
               text=[f"{judge_vals[d]:.3f}" for d in datasets], textposition="outside"),
    ])
    fig.update_layout(
        title="Same chatbot, two verdicts — surface metrics vs the LLM judge",
        barmode="group", yaxis={"range": [0, 1.1], "title": "Score (0–1, higher = better)"},
        paper_bgcolor=_BG, plot_bgcolor=_BG,
        legend={"orientation": "h", "yanchor": "bottom", "y": 1.02, "xanchor": "right", "x": 1},
        margin={"l": 50, "r": 20, "t": 80, "b": 50}, height=380,
    )
    return fig


def chart_judge_distribution() -> go.Figure:
    _, judge = _load_summary_and_judge()
    if not judge or not judge.get("rows"):
        return _placeholder_figure("Run the Gemini judge first (step 6 in the README)")

    scores = [r["score"] for r in judge["rows"]]
    fig = go.Figure(data=[
        go.Histogram(x=scores, nbinsx=11, marker_color=_JUDGE_COLOR,
                     hovertemplate="score bucket: %{x}<br>count: %{y}<extra></extra>"),
    ])
    n_zero = sum(1 for s in scores if s == 0.0)
    n_one = sum(1 for s in scores if s == 1.0)
    fig.update_layout(
        title=f"Judge score distribution (n={len(scores)}, mean={judge.get('mean_score', 0):.3f})  ·  "
              f"bimodal: {n_zero} hard-0s, {n_one} clean-1s",
        xaxis={"title": "judge score (0 = wrong, 1 = correct)", "range": [-0.05, 1.05]},
        yaxis={"title": "count of prompts"},
        paper_bgcolor=_BG, plot_bgcolor=_BG,
        margin={"l": 50, "r": 20, "t": 80, "b": 50}, height=380,
        showlegend=False,
    )
    return fig


def latest_results_summary() -> tuple[str, str]:
    rd = _latest_run_dir()
    if not rd:
        return ("_No CeRAI run results checked into the repo yet. "
                "Run `bash scripts/run_full_eval.sh` locally to populate `results/`._",
                "{}")
    summary = json.loads((rd / "summary.json").read_text())
    md_lines = [f"### Run `{summary['run_name']}`",
                f"- Total testcases: **{summary['total_testcases']}**",
                f"- Datasets: " + ", ".join(
                    f"{k} ({v})" for k, v in summary.get("datasets", {}).items()),
                "",
                "| Dataset | Metric | n | Mean score | Mean latency (s) | MCQ accuracy |",
                "|---|---|---|---|---|---|"]
    for ds, by_metric in summary.get("by_dataset_metric", {}).items():
        for metric, info in by_metric.items():
            mcq = info.get("mcq")
            mcq_text = (
                "—" if not mcq else
                f"{mcq['correct']}/{mcq['n']} ({mcq['accuracy']*100:.0f}%)")
            md_lines.append(
                f"| {ds.upper()} | {metric} | {info['n']} | "
                f"{info['score_mean'] if info['score_mean'] is not None else '—'} | "
                f"{info['latency_mean'] if info['latency_mean'] is not None else '—'} | "
                f"{mcq_text} |"
            )
    md_lines.append("")
    md_lines.append(f"_Per-prompt detail: see `results/{summary['run_name']}/results.md` "
                    f"in the repo._")
    return "\n".join(md_lines), json.dumps(summary, indent=2)


# --------------------------------------------------------------------------- #
# Live demo handlers                                                           #
# --------------------------------------------------------------------------- #
def amul_chat(message, history):  # untyped on purpose — list[dict] tripped
    # gradio_client's JSON-schema → Python-type resolver on the /api/info route
    # (it can't handle additionalProperties: true under dict-without-parameters).
    if not message or not str(message).strip():
        return "Type a question first."
    text, elapsed = call_amul(str(message))
    footer = f"\n\n_(latency: {elapsed:.2f}s · session: fresh)_"
    return text + footer


# --------------------------------------------------------------------------- #
# Mini-eval (in-browser, deterministic, small)                                 #
# --------------------------------------------------------------------------- #
MCQ_LETTER_RE = re.compile(r"\s*\*?\*?([A-D])[\.\)\*\s]?")

MINI_EVAL = [
    {
        "q": "The normal titratable acidity in fresh cow milk is approximately 0.13 to 0.17%. True or false?",
        "expected_substrings": ["true", "yes", "0.1", "correct"],
        "kind": "free_form",
    },
    {
        "q": ("Question: In HTST pasteurization milk is heated to at least?\n"
              "A) 63 °C\nB) 71.5 °C\nC) 75.5 °C\nD) 80 °C\n\n"
              "Reply with ONLY the single correct letter."),
        "expected_letter": "B",
        "kind": "mcq",
    },
    {
        "q": ("Question: As per PFA rules fat content in skim milk should be?\n"
              "A) Not less than 0.25%\nB) Not more than 0.25%\n"
              "C) Not more than 0.5 %\nD) Not less than 0.5 %\n\n"
              "Reply with ONLY the single correct letter."),
        "expected_letter": "C",
        "kind": "mcq",
    },
    {
        "q": "What is the typical lactation yield of a Sahiwal cow in litres per lactation?",
        "expected_substrings": ["sahiwal", "kg", "litre", "1500", "2000", "1700", "1800"],
        "kind": "free_form",
    },
]


def run_mini_eval(progress=gr.Progress()) -> tuple[str, str]:
    rows = []
    correct = 0
    total = 0
    for i, item in enumerate(MINI_EVAL):
        progress(i / len(MINI_EVAL), desc=f"Asking Amul AI (q {i+1}/{len(MINI_EVAL)})")
        try:
            text, elapsed = call_amul(item["q"])
        except httpx.HTTPStatusError as exc:
            rows.append((item["q"][:60], "ERROR",
                         f"upstream {exc.response.status_code}: {exc.response.text[:100]}",
                         "—", "—"))
            continue
        except httpx.RequestError as exc:
            rows.append((item["q"][:60], "ERROR",
                         f"network error: {type(exc).__name__}: {str(exc)[:120]}",
                         "—", "—"))
            continue
        except Exception as exc:
            rows.append((item["q"][:60], "ERROR",
                         f"{type(exc).__name__}: {str(exc)[:120]}", "—", "—"))
            continue
        if item["kind"] == "mcq":
            m = MCQ_LETTER_RE.match(text.strip())
            got = m.group(1).upper() if m else "?"
            ok = got == item["expected_letter"]
            correct += int(ok)
            total += 1
            rows.append((item["q"][:60], item["expected_letter"], got,
                         f"{elapsed:.2f}s", "✅" if ok else "❌"))
        else:
            blob = text.lower()
            ok = any(s in blob for s in item["expected_substrings"])
            rows.append((item["q"][:60], "/".join(item["expected_substrings"]),
                         text[:80], f"{elapsed:.2f}s", "✅" if ok else "❔"))
    progress(1.0, desc="Done")
    score_text = (f"**Score**: {correct}/{total} MCQ correct" if total
                  else "**Score**: free-form rows are flagged with ✅ if any expected substring appears.")
    md = ["| Prompt (truncated) | Expected | Got | Latency | OK |",
          "|---|---|---|---|---|"]
    for r in rows:
        md.append(f"| {r[0]} | {r[1]} | {r[2]} | {r[3]} | {r[4]} |")
    return score_text, "\n".join(md)


# --------------------------------------------------------------------------- #
# UI                                                                           #
# --------------------------------------------------------------------------- #
HEADER = """
# Evaluating Amul AI with CeRAI's *AIEvaluationTool*


> Started Option A (use the tool to evaluate a real conversational endpoint).
> Found enough fundamental issues mid-run that the writeup pivots into Option B
> territory, empirical run, concrete extension proposal grounded in
> 2025/26 evaluation literature. Both halves of evidence are in this Space and
> the linked GitHub repo.
"""


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="Amul AI × CeRAI — Eval submission",
                   theme=gr.themes.Soft()) as demo:
        gr.Markdown(HEADER)

        with gr.Tabs():
            with gr.Tab("Results at a glance"):
                gr.Markdown(
                    "## Headline numbers from the latest CeRAI run\n\n"
                    "Three plots that summarise what this evaluation found. The "
                    "narrative version is in the **Findings** tab; the per-prompt "
                    "table is in **CeRAI run results**.\n\n"
                    "**Why the chatbot looks bad here and good in the judge plot below "
                    "is the entire argument of this submission** — surface-overlap "
                    "metrics (BLEU/ROUGE/METEOR) penalise long English responses "
                    "against short or Romanised-Hindi references; an LLM judge with "
                    "an escape-hatch rubric tells a very different (and arguably more "
                    "honest) story about the same data."
                )
                gr.Plot(chart_surface_metrics_per_dataset())
                gr.Plot(chart_surface_vs_judge())
                gr.Plot(chart_judge_distribution())
                with gr.Accordion("How to read these", open=False):
                    gr.Markdown(
                        "- **Chart 1**: CeRAI's surface metrics on each dataset. "
                        "BBK is short option-text references (ROUGE picks up some "
                        "overlap, BLEU/METEOR don't); KCC is Romanised-Hindi short "
                        "references vs the chatbot's English replies — overlap "
                        "collapses to ~0.\n"
                        "- **Chart 2**: Same chatbot, two verdicts. The surface "
                        "metrics average ≈ 0.30 on BBK and ≈ 0.01 on KCC. The LLM "
                        "judge (with the rubric in `extension_proposal.md` §1/§3) "
                        "scores them ≈ 0.79 and ≈ 1.00. That gap is the headline.\n"
                        "- **Chart 3**: The judge is bimodal — most prompts are clean "
                        "1.0s, a small set are clean 0.0s (factual misses), only a "
                        "couple sit in the middle. That shape is more honest than "
                        "the surface metrics' diffuse 0.0–0.4 cloud."
                    )

            for label, fname, fallback in REPORT_DOCS:
                with gr.Tab(label):
                    gr.Markdown(_read_or_placeholder(REPORT_DIR / fname, fallback))

            with gr.Tab("Try Amul AI"):
                gr.Markdown(
                    "Calls `https://api.prod.amulai.in/api/chat/` with an anonymous "
                    "JWT — same code path as our `amul_proxy/`. Each turn is a fresh "
                    "session (no chat memory)."
                )
                gr.ChatInterface(
                    fn=amul_chat,
                    type="messages",
                    cache_examples=False,  # examples cache builds JSON schema -> crash on bool
                    examples=[
                        "How many liters of milk does a Holstein cow produce daily?",
                        "What is the normal titratable acidity of fresh cow milk?",
                        "Pasu ko khurpaka munhpaka ki bimari ho gayi hai, kya karu?",
                        "How do I prevent mastitis in my buffalo?",
                    ],
                )

            with gr.Tab("CeRAI run results"):
                md, raw = latest_results_summary()
                gr.Markdown(md)
                with gr.Accordion("Raw summary.json", open=False):
                    gr.Code(value=raw, language="json")

            with gr.Tab("Run mini-eval"):
                gr.Markdown(
                    "Hits Amul AI with a tiny built-in 4-question slice (2 MCQ, "
                    "2 free-form) drawn from BhashaBench-Krishi dairy/poultry. "
                    "Useful for a quick smoke test from the browser; the *real* "
                    "evaluation is in the **CeRAI run results** tab."
                )
                run_btn = gr.Button("Run mini-eval", variant="primary")
                score_md = gr.Markdown()
                table_md = gr.Markdown()
                run_btn.click(run_mini_eval, outputs=[score_md, table_md])

            with gr.Tab("How to reproduce"):
                gr.Markdown(_read_or_placeholder(
                    ROOT / "README.md",
                    "_README missing._"))

    return demo


if __name__ == "__main__":
    demo = build_ui()
    demo.queue(default_concurrency_limit=2)
    demo.launch(server_name="0.0.0.0",
                server_port=int(os.environ.get("PORT", 7860)),
                show_api=False)  # avoid /api/info schema crash in gradio_client 1.14