#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import json
import re
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any

import requests

from src.rag.project_config import env


CONFIGS = [
    "A0_full",
    "A1_no_planner",
    "A2_no_neo4j",
    "A3_no_vector",
    "A4_one_hop",
    "A5_no_bge",
]


def atomic_json(path: Path, payload: Any) -> None:
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    tmp.replace(path)


def extract_json(text: str) -> dict[str, Any]:
    text = (text or "").strip()
    try:
        return json.loads(text)
    except Exception:
        pass
    fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.S)
    if fenced:
        try:
            return json.loads(fenced.group(1))
        except Exception:
            pass
    match = re.search(r"\{.*\}", text, re.S)
    if match:
        try:
            return json.loads(match.group(0))
        except Exception:
            pass
    return {}


def compact_sources(sources: list[dict[str, Any]]) -> list[dict[str, Any]]:
    compact: list[dict[str, Any]] = []
    seen: set[str] = set()
    for source in sources:
        evidence = " ".join(str(source.get("evidence") or "").split())
        key = evidence.lower()
        if not evidence or key in seen:
            continue
        seen.add(key)
        compact.append(
            {
                "srt": (
                    f"{source.get('source','')} -- "
                    f"{source.get('relation','')} --> "
                    f"{source.get('target','')}"
                ),
                "evidence": evidence[:260],
            }
        )
        if len(compact) >= 5:
            break
    return compact


def judge_one(
    *,
    root: Path,
    output_dir: Path,
    gold: dict[str, Any],
    api_url: str,
    api_key: str,
    model: str,
    retries: int,
) -> dict[str, Any]:
    qid = str(gold["id"])
    output_path = output_dir / f"{qid}.json"
    if output_path.exists():
        try:
            existing = json.loads(output_path.read_text(encoding="utf-8"))
            judgments = existing.get("judgments") or []
            if len(judgments) == len(CONFIGS):
                return {"question_id": qid, "status": "ok", "resumed": True}
        except Exception:
            pass

    candidates: list[dict[str, Any]] = []
    for config in CONFIGS:
        record = json.loads(
            (root / config / f"{qid}.json").read_text(encoding="utf-8")
        )
        response = record.get("response") or {}
        candidates.append(
            {
                "config_id": config,
                "answer": response.get("answer", ""),
                "retrieved_evidence": compact_sources(response.get("sources") or []),
            }
        )

    prompt = f"""
You are evaluating six RAG answers to one plant-biology question.

The GOLD SOURCE TEXT is authoritative. The short GOLD ANSWER may be incomplete;
if they differ, judge against the GOLD SOURCE TEXT.
Do not reward verbosity. Do not penalize a concise correct answer.
Judge each system independently.

Question:
{gold['question']}

Gold answer:
{gold['gold_answer']}

Gold source text:
{gold['source_text']}

Candidate systems:
{json.dumps(candidates, ensure_ascii=False)}

For each configuration, judge:
1. answer_verdict:
   - correct: all facts needed by the question are correct; extra claims do not conflict.
   - partial: some required facts are correct but important requested facts are missing,
     or the answer mixes the target result with conflicting alternatives.
   - wrong: gives a conflicting or unrelated substantive answer.
   - insufficient: explicitly cannot answer or gives no substantive answer.
2. evidence_verdict:
   - full: retrieved evidence contains all facts needed to answer.
   - partial: retrieved evidence contains some but not all required facts.
   - miss: retrieved evidence does not contain the target fact.
3. usable: true only when answer_verdict is correct or partial.

Return JSON only:
{{
  "question_id": "{qid}",
  "judgments": [
    {{
      "config_id": "A0_full",
      "answer_verdict": "correct|partial|wrong|insufficient",
      "evidence_verdict": "full|partial|miss",
      "usable": true,
      "reason": "maximum 15 words"
    }}
  ]
}}
""".strip()

    last_error = ""
    raw = ""
    for attempt in range(1, retries + 2):
        try:
            response = requests.post(
                api_url,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": model,
                    "messages": [
                        {
                            "role": "system",
                            "content": "Return valid JSON only. Apply the evaluation rubric consistently.",
                        },
                        {"role": "user", "content": prompt},
                    ],
                    "temperature": 0,
                    "max_tokens": 4000,
                },
                timeout=240,
            )
            response.raise_for_status()
            data = response.json()
            raw = (data.get("choices") or [{}])[0].get("message", {}).get(
                "content", ""
            )
            parsed = extract_json(raw)
            judgments = parsed.get("judgments") or []
            config_ids = {str(item.get("config_id")) for item in judgments}
            if len(judgments) == len(CONFIGS) and config_ids == set(CONFIGS):
                payload = {
                    "question_id": qid,
                    "question": gold["question"],
                    "gold_answer": gold["gold_answer"],
                    "gold_source_text": gold["source_text"],
                    "judgments": judgments,
                    "attempt": attempt,
                    "model": model,
                }
                atomic_json(output_path, payload)
                return {"question_id": qid, "status": "ok", "attempt": attempt}
            last_error = "Incomplete or invalid judgment JSON"
        except Exception as exc:
            last_error = f"{type(exc).__name__}: {exc}"
        time.sleep(min(4 * attempt, 12))

    atomic_json(
        output_path,
        {
            "question_id": qid,
            "question": gold["question"],
            "error": last_error,
            "raw": raw,
            "model": model,
        },
    )
    return {"question_id": qid, "status": "error", "error": last_error}


def summarize(
    root: Path,
    judge_dir: Path,
    questions: list[dict[str, Any]],
    output_dir: Path,
) -> None:
    rows: list[dict[str, Any]] = []
    for gold in questions:
        qid = str(gold["id"])
        judged = json.loads((judge_dir / f"{qid}.json").read_text(encoding="utf-8"))
        for item in judged.get("judgments") or []:
            config = str(item["config_id"])
            record = json.loads(
                (root / config / f"{qid}.json").read_text(encoding="utf-8")
            )
            response = record.get("response") or {}
            rows.append(
                {
                    "question_id": qid,
                    "config_id": config,
                    "question": gold["question"],
                    "gold_answer": gold["gold_answer"],
                    "answer_verdict": item.get("answer_verdict", ""),
                    "evidence_verdict": item.get("evidence_verdict", ""),
                    "usable": bool(item.get("usable")),
                    "reason": item.get("reason", ""),
                    "answer": response.get("answer", ""),
                    "source_count": len(response.get("sources") or []),
                    "result_file": str(root / config / f"{qid}.json"),
                }
            )
    csv_path = output_dir / "llm_judgments_100q.csv"
    with csv_path.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)

    summaries: list[dict[str, Any]] = []
    for config in CONFIGS:
        subset = [row for row in rows if row["config_id"] == config]
        answers = Counter(str(row["answer_verdict"]) for row in subset)
        evidence = Counter(str(row["evidence_verdict"]) for row in subset)
        n = len(subset)
        weighted_answer = (
            answers["correct"] + 0.5 * answers["partial"]
        ) / n
        weighted_recall = (
            evidence["full"] + 0.5 * evidence["partial"]
        ) / n
        summaries.append(
            {
                "config_id": config,
                "questions": n,
                "correct": answers["correct"],
                "partial": answers["partial"],
                "wrong": answers["wrong"],
                "insufficient": answers["insufficient"],
                "strict_answer_accuracy": answers["correct"] / n,
                "weighted_answer_accuracy": weighted_answer,
                "usable_answer_rate": sum(
                    1 for row in subset if row["usable"]
                )
                / n,
                "full_evidence_hits": evidence["full"],
                "partial_evidence_hits": evidence["partial"],
                "evidence_misses": evidence["miss"],
                "question_level_full_recall": evidence["full"] / n,
                "question_level_weighted_recall": weighted_recall,
            }
        )
    summary = {
        "evaluation": "DeepSeek LLM-assisted judge",
        "warning": (
            "These are automated judge estimates. They are not a substitute for "
            "human review and should be calibrated against the manually reviewed 20 questions."
        ),
        "configs": summaries,
    }
    atomic_json(output_dir / "llm_judgment_summary_100q.json", summary)
    with (output_dir / "llm_judgment_summary_100q.csv").open(
        "w", encoding="utf-8-sig", newline=""
    ) as handle:
        writer = csv.DictWriter(handle, fieldnames=list(summaries[0]))
        writer.writeheader()
        writer.writerows(summaries)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-dir", required=True)
    parser.add_argument("--questions", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--retries", type=int, default=2)
    args = parser.parse_args()

    api_key = env("DEEPSEEK_API_KEY")
    if not api_key:
        raise RuntimeError("No DEEPSEEK_API_KEY configured")
    base_url = env("DEEPSEEK_BASE_URL", "https://api.deepseek.com").rstrip("/")
    model = env("V15_RELATION_LLM_MODEL", "deepseek-v4-flash")
    questions = json.loads(Path(args.questions).read_text(encoding="utf-8-sig"))
    root = Path(args.input_dir)
    output_dir = Path(args.output_dir)
    judge_dir = output_dir / "per_question"
    judge_dir.mkdir(parents=True, exist_ok=True)
    progress_path = output_dir / "judge_progress.json"

    completed: list[dict[str, Any]] = []
    with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool:
        futures = [
            pool.submit(
                judge_one,
                root=root,
                output_dir=judge_dir,
                gold=gold,
                api_url=f"{base_url}/chat/completions",
                api_key=api_key,
                model=model,
                retries=args.retries,
            )
            for gold in questions
        ]
        for index, future in enumerate(as_completed(futures), 1):
            result = future.result()
            completed.append(result)
            atomic_json(
                progress_path,
                {
                    "completed": index,
                    "total": len(questions),
                    "errors": [x for x in completed if x["status"] == "error"],
                },
            )
            print(
                f"[{index}/{len(questions)}] {result['question_id']} "
                f"{result['status']}",
                flush=True,
            )
    errors = [item for item in completed if item["status"] == "error"]
    if errors:
        raise RuntimeError(f"Judge failures: {errors}")
    summarize(root, judge_dir, questions, output_dir)


if __name__ == "__main__":
    main()