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

import argparse
import csv
import hashlib
import html
import json
import re
import unicodedata
from collections import Counter
from difflib import SequenceMatcher
from pathlib import Path
from typing import Any

from ir_measures import P, RR, Recall, nDCG, Qrel, ScoredDoc, calc_aggregate


CONFIGS = [
    "A0_full",
    "A1_no_planner",
    "A2_no_neo4j",
    "A3_no_vector",
    "A4_one_hop",
    "A5_no_bge",
]
KS = (5, 10, 30)


def normalize(text: Any) -> str:
    value = html.unescape(str(text or ""))
    value = unicodedata.normalize("NFKC", value).lower()
    value = re.sub(r"<[^>]+>", " ", value)
    value = re.sub(r"https?://\S+", " ", value)
    value = re.sub(r"[^a-z0-9\u4e00-\u9fff.%+-]+", " ", value)
    return " ".join(value.split())


def tokens(text: Any) -> list[str]:
    return re.findall(r"[a-z0-9.%+-]+|[\u4e00-\u9fff]", normalize(text))


def evidence_similarity(gold: str, candidate: str) -> dict[str, float | bool]:
    ng = normalize(gold)
    nc = normalize(candidate)
    if not ng or not nc:
        return {
            "matched": False,
            "containment": False,
            "token_recall": 0.0,
            "token_precision": 0.0,
            "sequence_ratio": 0.0,
            "match_score": 0.0,
        }
    containment = ng in nc or nc in ng
    gt = set(tokens(gold))
    ct = set(tokens(candidate))
    overlap = len(gt & ct)
    token_recall = overlap / len(gt) if gt else 0.0
    token_precision = overlap / len(ct) if ct else 0.0
    sequence_ratio = SequenceMatcher(None, ng, nc).ratio()
    match_score = max(
        sequence_ratio,
        0.72 * token_recall + 0.28 * token_precision,
    )
    matched = bool(
        containment
        or sequence_ratio >= 0.82
        or (token_recall >= 0.72 and token_precision >= 0.28)
    )
    return {
        "matched": matched,
        "containment": containment,
        "token_recall": token_recall,
        "token_precision": token_precision,
        "sequence_ratio": sequence_ratio,
        "match_score": match_score,
    }


def evidence_doc_id(evidence: str) -> str:
    digest = hashlib.sha1(normalize(evidence).encode("utf-8")).hexdigest()[:20]
    return f"EVID::{digest}"


def paper_id(source: dict[str, Any]) -> str:
    return str(
        source.get("paper_id")
        or source.get("pmcid")
        or source.get("pmid")
        or source.get("doi")
        or source.get("paper_title")
        or source.get("title")
        or ""
    ).strip()


def dedupe_ranked_sources(
    qid: str,
    gold_source: str,
    sources: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    ranked: list[dict[str, Any]] = []
    audit: list[dict[str, Any]] = []
    seen_evidence: set[str] = set()
    gold_doc_id = f"GOLD::{qid}"
    for raw_rank, source in enumerate(sources, 1):
        evidence = str(source.get("evidence") or "")
        normalized = normalize(evidence)
        if not normalized or normalized in seen_evidence:
            continue
        seen_evidence.add(normalized)
        similarity = evidence_similarity(gold_source, evidence)
        doc_id = gold_doc_id if similarity["matched"] else evidence_doc_id(evidence)
        row = {
            "rank": len(ranked) + 1,
            "raw_rank": raw_rank,
            "doc_id": doc_id,
            "paper_id": paper_id(source),
            "evidence": evidence,
            "source": source.get("source", ""),
            "relation": source.get("relation", ""),
            "target": source.get("target", ""),
            **similarity,
        }
        ranked.append(row)
        audit.append(row)
    return ranked, audit


def metric_key(measure: Any) -> str:
    return str(measure)


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    with path.open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--results-dir", required=True)
    parser.add_argument("--questions", required=True)
    parser.add_argument("--judgments", required=True)
    parser.add_argument("--output-dir", required=True)
    args = parser.parse_args()

    results_dir = Path(args.results_dir)
    questions = json.loads(Path(args.questions).read_text(encoding="utf-8-sig"))
    question_by_id = {str(item["id"]): item for item in questions}
    judgment_rows = list(
        csv.DictReader(open(args.judgments, encoding="utf-8-sig"))
    )
    judgment_by_key = {
        (row["config_id"], row["question_id"]): row for row in judgment_rows
    }
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "trec_runs").mkdir(exist_ok=True)

    qrels = [
        Qrel(str(item["id"]), f"GOLD::{item['id']}", 1) for item in questions
    ]
    qrels_path = output_dir / "known_item_qrels.tsv"
    with qrels_path.open("w", encoding="utf-8", newline="") as handle:
        for item in questions:
            qid = str(item["id"])
            handle.write(f"{qid}\t0\tGOLD::{qid}\t1\n")

    measures = [
        P @ 5,
        P @ 10,
        P @ 30,
        Recall @ 5,
        Recall @ 10,
        Recall @ 30,
        RR @ 30,
        nDCG @ 5,
        nDCG @ 10,
        nDCG @ 30,
    ]
    summary_rows: list[dict[str, Any]] = []
    query_rows: list[dict[str, Any]] = []
    match_audit: list[dict[str, Any]] = []

    for config in CONFIGS:
        run: list[ScoredDoc] = []
        diversity: dict[int, list[float]] = {k: [] for k in KS}
        diversity_returned: dict[int, list[float]] = {k: [] for k in KS}
        duplicate_rates: list[float] = []
        retrieved_counts: list[int] = []
        per_question_ranked: dict[str, list[dict[str, Any]]] = {}

        run_path = output_dir / "trec_runs" / f"{config}.run"
        with run_path.open("w", encoding="utf-8", newline="") as run_handle:
            for item in questions:
                qid = str(item["id"])
                record = json.loads(
                    (results_dir / config / f"{qid}.json").read_text(
                        encoding="utf-8"
                    )
                )
                response = record.get("response") or {}
                raw_sources = response.get("sources") or []
                ranked, audit = dedupe_ranked_sources(
                    qid, str(item.get("source_text") or ""), raw_sources
                )
                per_question_ranked[qid] = ranked
                retrieved_counts.append(len(ranked))
                duplicate_rates.append(
                    1 - (len(ranked) / len(raw_sources)) if raw_sources else 0.0
                )
                for row in audit:
                    match_audit.append(
                        {
                            "config_id": config,
                            "question_id": qid,
                            "gold_source_text": item.get("source_text", ""),
                            **row,
                        }
                    )
                for row in ranked:
                    score = 10000.0 - float(row["rank"])
                    run.append(ScoredDoc(qid, str(row["doc_id"]), score))
                    run_handle.write(
                        f"{qid} Q0 {row['doc_id']} {row['rank']} "
                        f"{score:.6f} {config}\n"
                    )
                for k in KS:
                    top = ranked[:k]
                    unique_papers = len(
                        {str(row["paper_id"]) for row in top if row["paper_id"]}
                    )
                    diversity[k].append(unique_papers / k)
                    diversity_returned[k].append(
                        unique_papers / len(top) if top else 0.0
                    )

        aggregate = calc_aggregate(measures, qrels, run)
        summary: dict[str, Any] = {
            "config_id": config,
            "questions": len(questions),
        }
        for measure in measures:
            summary[metric_key(measure)] = float(aggregate[measure])
        for k in KS:
            summary[f"source_diversity@{k}"] = sum(diversity[k]) / len(
                diversity[k]
            )
            summary[f"source_diversity_returned@{k}"] = sum(
                diversity_returned[k]
            ) / len(diversity_returned[k])
        summary["mean_exact_duplicate_rate"] = sum(duplicate_rates) / len(
            duplicate_rates
        )
        summary["mean_deduplicated_source_count"] = sum(retrieved_counts) / len(
            retrieved_counts
        )

        answer_counts = Counter()
        answer_scores: list[float] = []
        for item in questions:
            qid = str(item["id"])
            judgment = judgment_by_key[(config, qid)]
            verdict = judgment["answer_verdict"]
            answer_counts[verdict] += 1
            answer_scores.append(
                {"correct": 1.0, "partial": 0.5}.get(verdict, 0.0)
            )
            ranked = per_question_ranked[qid]
            first_rank = next(
                (
                    int(row["rank"])
                    for row in ranked
                    if row["doc_id"] == f"GOLD::{qid}"
                ),
                0,
            )
            query_rows.append(
                {
                    "config_id": config,
                    "question_id": qid,
                    "question": item["question"],
                    "answer_verdict": verdict,
                    "answer_score": {
                        "correct": 1.0,
                        "partial": 0.5,
                    }.get(verdict, 0.0),
                    "evidence_verdict_llm": judgment["evidence_verdict"],
                    "known_gold_first_rank": first_rank,
                    "known_gold_hit@5": int(bool(first_rank and first_rank <= 5)),
                    "known_gold_hit@10": int(
                        bool(first_rank and first_rank <= 10)
                    ),
                    "known_gold_hit@30": int(
                        bool(first_rank and first_rank <= 30)
                    ),
                    "deduplicated_source_count": len(ranked),
                    "unique_papers@10": len(
                        {
                            str(row["paper_id"])
                            for row in ranked[:10]
                            if row["paper_id"]
                        }
                    ),
                    "judge_reason": judgment["reason"],
                }
            )
        summary["answer_correct"] = answer_counts["correct"]
        summary["answer_partial"] = answer_counts["partial"]
        summary["answer_wrong"] = answer_counts["wrong"]
        summary["answer_insufficient"] = answer_counts["insufficient"]
        summary["answer_correctness_1_0.5_0"] = sum(answer_scores) / len(
            answer_scores
        )
        summary_rows.append(summary)

        scored_rows: list[dict[str, Any]] = []
        for item in questions:
            qid = str(item["id"])
            judgment = judgment_by_key[(config, qid)]
            verdict = judgment["answer_verdict"]
            record = json.loads(
                (results_dir / config / f"{qid}.json").read_text(
                    encoding="utf-8"
                )
            )
            scored_rows.append(
                {
                    "id": qid,
                    "question": item["question"],
                    "gold_answer": item["gold_answer"],
                    "deepseek_answer": (record.get("response") or {}).get(
                        "answer", ""
                    ),
                    "score": {"correct": 1.0, "partial": 0.5}.get(
                        verdict, 0.0
                    ),
                    "label": verdict,
                    "reason": judgment["reason"],
                    "model": "DeepSeek V4Flash + RAG",
                    "status": "answered",
                }
            )
        config_base = output_dir / f"{config}_same_scale_correctness"
        (config_base.with_suffix(".json")).write_text(
            json.dumps(scored_rows, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        write_csv(config_base.with_suffix(".csv"), scored_rows)
        config_summary = {
            "model": f"DeepSeek V4Flash + RAG ({config})",
            "n": len(scored_rows),
            "total_score": sum(float(row["score"]) for row in scored_rows),
            "answer_correctness": summary["answer_correctness_1_0.5_0"],
            "label_counts": dict(answer_counts),
            "score_counts": dict(
                Counter(str(row["score"]) for row in scored_rows)
            ),
            "scoring_note": (
                "Unified 1/0.5/0 scale. Labels come from the calibrated "
                "DeepSeek evidence-aware judge, not hard-coded question IDs."
            ),
        }
        (output_dir / f"{config}_same_scale_correctness_summary.json").write_text(
            json.dumps(config_summary, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    write_csv(output_dir / "ir_metrics_summary.csv", summary_rows)
    write_csv(output_dir / "ir_metrics_per_question.csv", query_rows)
    write_csv(output_dir / "gold_evidence_match_audit.csv", match_audit)
    report = {
        "protocol": "known-item retrieval evaluation",
        "implementation": "ir_measures 0.4.1 / pytrec-eval-terrier",
        "qrels": str(qrels_path),
        "ks": KS,
        "matching_rule": {
            "containment": True,
            "sequence_ratio_threshold": 0.82,
            "token_recall_threshold": 0.72,
            "token_precision_threshold": 0.28,
        },
        "important_limitation": (
            "Each question has one known gold source text. Therefore Recall@k "
            "is equivalent to known-item hit rate. Precision@k and nDCG@k "
            "treat unjudged evidence as non-relevant and may underestimate "
            "true retrieval quality. Formal publication-grade P/nDCG requires "
            "pooled relevance judgments over retrieved evidence."
        ),
        "summary": summary_rows,
    }
    (output_dir / "ir_evaluation_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()