"""Measure generated answers on a declared, frozen graph benchmark.

Regression mode repeats known questions; it is never a new independent TEST.
The harness supplies retrieval context and does not test native tool routing.
"""

from __future__ import annotations

import argparse
import concurrent.futures
import csv
import json
import os
import time
from collections import Counter
from pathlib import Path

from tools.graphrag.graph_db import create_neo4j_driver, get_neo4j_settings
from tools.graphrag.graph_search import search_graph_with_session
from tools.graphrag.build_graphrag import graph_version
from tools.graphrag.evaluation.check_model_sample import grade_claims
from tools.graphrag.evaluation.reproducibility import (
    graph_fingerprint,
    run_metadata,
    sha256_file,
    verify_freeze,
)
from tools.graphrag.evaluation.source_oracle import load_source_oracle, norm
from tools.rag.evaluation.rag_runner import DEFAULT_MODEL, OPENWEBUI_URL, request_json
from tools.answers.checked_generation import generate_checked_answer
from tools.answers.graph_context import graph_context
from tools.answers.context_policy import (
    citation_check,
    allowed_urls,
    answer_context_issues,
    prose_issues,
    text_contains_value,
)

HERE = Path(__file__).resolve().parents[1] / "answers"


def grade_answer(case: dict, answer: dict, oracle: dict, *, context: dict | None = None) -> dict:
    # JSON null is absence, not a positive factual assertion. The same rule is
    # applicable to all rounds, including the archived initial sample.
    normalized = dict(answer)
    if isinstance(answer.get("claims"), list):
        normalized["claims"] = [
            {k: v for k, v in claim.items() if v is not None} if isinstance(claim, dict) else claim
            for claim in answer["claims"]
        ]
    grade = grade_claims(case, normalized, oracle)
    urls = (
        allowed_urls(context)
        if context is not None
        else {d["source_url"] for d in oracle.get("documents", {}).values() if d.get("source_url")}
    )
    citations = citation_check(answer, urls)
    grade.update(citations)
    policy_issues = answer_context_issues(normalized, context) if context is not None else []
    grade["context_issues"] = policy_issues
    text = norm(answer.get("answer", ""))
    people = {str(b["person"]) for b in case["graph"]["expected_bindings"] if b.get("person")}
    missing = [
        v
        for v in case.get("expected_answer_contains", [])
        if not text_contains_value(answer.get("answer", ""), str(v), person=str(v) in people)
    ]
    policy_issues.extend(prose_issues(normalized))
    # The initial sample exposed a real contradiction in otherwise correct
    # category claims. Check it explicitly rather than silently accepting prose.
    wrong_thesis_label = (
        "dizertac" in text
        and any(
            str(b.get("category", "")).startswith("dp") for b in case["graph"]["expected_bindings"]
        )
        and not any(
            str(b.get("category", "")).startswith("phd") for b in case["graph"]["expected_bindings"]
        )
    )
    grade.update(
        text_values_ok=not missing,
        missing_text_values=missing,
        contradictory_thesis_label=wrong_thesis_label,
    )
    grade["strict"] = bool(
        grade["strict"]
        and not missing
        and not wrong_thesis_label
        and citations["citations_ok"]
        and not policy_issues
    )
    return grade


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--workers", type=int, default=3)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--split", choices=["dev", "test", "regression"], default="test")
    parser.add_argument("--questions", type=Path, required=True)
    parser.add_argument("--audit", type=Path, required=True)
    parser.add_argument("--freeze-manifest", type=Path, required=True)
    parser.add_argument("--dev-ids", type=Path, required=True)
    parser.add_argument("--regression-ids", type=Path)
    parser.add_argument(
        "--dev-results", type=Path, help="Required completed DEV before final regression."
    )
    args = parser.parse_args()
    if args.output_dir.exists():
        parser.error("Refusing to overwrite an answer-level TEST.")
    if not os.getenv("OPENWEBUI_API_KEY"):
        parser.error("OPENWEBUI_API_KEY missing.")
    if args.workers < 1:
        parser.error("--workers must be positive.")
    if args.split in {"regression", "test"} and args.dev_results is None:
        parser.error("Final regression requires completed --dev-results.")
    dataset = args.questions
    audit = args.audit
    frozen = args.freeze_manifest
    manifest = json.loads(frozen.read_text())
    questions = json.loads(dataset.read_text())
    oracle = load_source_oracle()
    development = set(json.loads(args.dev_ids.read_text()))
    test = [q for q in questions if q["split"] == "test"]
    if args.split == "regression":
        if args.regression_ids is None:
            parser.error("Regression needs explicit --regression-ids.")
        regression_ids = set(json.loads(args.regression_ids.read_text()))
        selected = [q for q in questions if q["id"] in regression_ids]
        if len(selected) != len(regression_ids):
            parser.error("Unknown/duplicate regression IDs.")
    else:
        selected = (
            test if args.split == "test" else [q for q in questions if q["id"] in development]
        )
    if not selected:
        parser.error("Empty selected split.")
    if args.split == "test" and (
        manifest.get("evaluation_kind") != "independent_holdout"
        or development & {q["id"] for q in selected}
    ):
        parser.error("Fresh TEST requires independent freeze and zero DEV overlap.")
    if not development <= {q["id"] for q in questions if q["split"] == "dev"}:
        parser.error("Generation DEV must be declared inside the dataset DEV split.")
    prompt_path = HERE / "graph_answer_prompt.txt"
    prompt = prompt_path.read_text()
    args.output_dir.mkdir(parents=True)
    metadata = run_metadata(dataset)
    settings = get_neo4j_settings()
    contexts = {}
    with (
        create_neo4j_driver(settings) as driver,
        driver.session(database=settings.database) as session,
    ):
        verify_freeze(
            manifest, metadata, graph_fingerprint(session), audit=audit, **manifest["parameters"]
        )
        metadata["index_version"] = graph_version(session)
        # Generate every context before making any external request.
        for case in selected:
            result = search_graph_with_session(session, case["question"], **manifest["parameters"])
            result["index_version"] = metadata["index_version"]
            contexts[case["id"]] = graph_context(result)
    code_files = [
        Path(__file__),
        prompt_path,
        HERE / "answer_format.py",
        HERE / "context_policy.py",
        HERE / "checked_generation.py",
        HERE / "render_answer.py",
        Path("tools/rag/evaluation/rag_runner.py"),
    ]
    hashes = {str(p.resolve().relative_to(HERE.parents[1])): sha256_file(p) for p in code_files}
    if args.split in {"regression", "test"}:
        dev = json.loads(args.dev_results.read_text())
        if (
            dev["configuration"]["selected_split"] != "dev"
            or set(dev["configuration"]["selected_ids"]) != development
            or dev["summary"]["questions"] != len(development)
            or {row["id"] for row in dev["questions"]} != development
            or len(dev["questions"]) != len(development)
            or dev["configuration"]["generation_files"] != hashes
            or dev["run_metadata"]["fingerprints"] != metadata["fingerprints"]
            or dev["configuration"]["graph_freeze_sha256"] != sha256_file(frozen)
            or dev["configuration"]["model"] != args.model
            or dev["run_metadata"].get("index_version") != metadata["index_version"]
        ):
            raise RuntimeError(
                "Generation code, prompt, data or complete DEV differs; rerun DEV before final freeze."
            )
    declaration = {
        "method": (
            "Known-question development selection."
            if args.split == "dev"
            else (
                "Previously exposed LLM questions; regression, not independent TEST."
                if args.split == "regression"
                else "Fresh query-group holdout over known sources and templates; zero DEV query-group overlap."
            )
        ),
        "graph_freeze_sha256": sha256_file(frozen),
        "audit_sha256": sha256_file(audit),
        "generation_files": hashes,
        "model": args.model,
        "temperature": 0,
        "workers": args.workers,
        "dev_ids": sorted(development),
        "test_ids": [q["id"] for q in test],
        "dev_count": len(development),
        "test_count": len(test),
        "target_accuracy": 0.92,
        "evaluation_kind": (
            "historical_regression"
            if args.split == "regression"
            else "independent_query_group_holdout" if args.split == "test" else "development"
        ),
        "rendering": "Verified structured claims rendered deterministically; raw model outputs preserved.",
        "selected_split": args.split,
        "selected_ids": [q["id"] for q in selected],
        "selected_count": len(selected),
        "selected_dev_overlap": len(development & {q["id"] for q in selected}),
        "corrective_request_limit": 1,
        "format_repair_limit": 1,
        "repair_policy": "At most one total corrective request for format, citations or retrieval-context violations. Same question/context, no gold; retain all attempts.",
        "dev_results_sha256": sha256_file(args.dev_results) if args.dev_results else None,
        "policy": (
            "Evaluate every declared regression case, including DEV overlap; retain all failures and stop tuning after this freeze."
            if args.split == "regression"
            else "Evaluate every declared fresh TEST case once; retain every failure; do not tune on this TEST."
        ),
        "selected_categories": dict(Counter(q["category"] for q in selected)),
        "limitations": [
            (
                "Known development questions; not a holdout."
                if args.split == "dev"
                else (
                    "All questions and previous model outputs are known; this is historical regression."
                    if args.split == "regression"
                    else "Known corpus/profiles and templates; these query groups were not previously evaluated."
                )
            ),
            "Harness supplies graph context; native OpenWebUI tool choice and browser UI are outside this experiment.",
            "Checks structured factual claims, their source support/citations and required prose values; not a general semantic judge of arbitrary prose.",
        ],
    }
    (args.output_dir / "frozen_answers.json").write_text(
        json.dumps(declaration, ensure_ascii=False, indent=2) + "\n"
    )
    results = {}

    def run(case):
        started = time.perf_counter()
        context = contexts[case["id"]]
        row = {
            "id": case["id"],
            "category": case["category"],
            "question": case["question"],
            "graph_context": context,
            "strict": False,
        }

        def complete(messages):
            return request_json(
                OPENWEBUI_URL,
                method="POST",
                timeout=90,
                max_attempts=2,
                headers={
                    "Authorization": f'Bearer {os.environ["OPENWEBUI_API_KEY"]}',
                    "Content-Type": "application/json",
                },
                payload={
                    "model": args.model,
                    "temperature": 0,
                    "stream": False,
                    "response_format": {"type": "json_object"},
                    "messages": messages,
                },
            )

        generated = generate_checked_answer(case["question"], context, prompt, complete)
        row.update(generated)
        # Grading occurs after all generation/correction decisions, never inside
        # the generator. Gold cannot influence whether another request is sent.
        for attempt in row["attempts"]:
            candidate = attempt.get("rendered_answer", attempt.get("answer"))
            attempt["strict"] = (
                grade_answer(case, candidate, oracle, context=context)["strict"]
                if candidate
                else False
            )
        if row.get("answer"):
            row.update(grade_answer(case, row["answer"], oracle, context=context))
        row["strict"] = bool(row["strict"] and row["accepted_for_delivery"])
        row["latency_seconds"] = time.perf_counter() - started
        return row

    output = {"configuration": declaration, "run_metadata": metadata, "questions": []}
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor:
        futures = [executor.submit(run, case) for case in selected]
        for future in concurrent.futures.as_completed(futures):
            row = future.result()
            results[row["id"]] = row
            output["questions"] = [results[q["id"]] for q in selected if q["id"] in results]
            output["summary"] = {
                "questions": len(results),
                "correct": sum(r["strict"] for r in results.values()),
                "accuracy": sum(r["strict"] for r in results.values()) / len(results),
                "errors": sum(bool(r.get("error")) for r in results.values()),
                "first_pass_correct": sum(
                    bool(r.get("attempts") and r["attempts"][0].get("strict"))
                    for r in results.values()
                ),
                "format_repair_attempted": sum(
                    bool(
                        r.get("attempts")
                        and r["attempts"][0].get("format_error")
                        and len(r["attempts"]) > 1
                    )
                    for r in results.values()
                ),
                "corrective_request_count": sum(
                    len(r.get("attempts", [])) > 1 for r in results.values()
                ),
                "context_repair_attempted": sum(
                    bool(
                        r.get("attempts")
                        and r["attempts"][0].get("context_issues")
                        and len(r["attempts"]) > 1
                    )
                    for r in results.values()
                ),
                "normalized_first_response": sum(
                    bool(r.get("attempts") and r["attempts"][0].get("format_normalized"))
                    for r in results.values()
                ),
            }
            (args.output_dir / f"answers_{args.split}.json").write_text(
                json.dumps(output, ensure_ascii=False, indent=2) + "\n"
            )
            print(
                f'{len(results)}/{len(selected)} {row["id"]}: {"OK" if row["strict"] else "FAIL"} {row["latency_seconds"]:.1f}s',
                flush=True,
            )
    with (
        create_neo4j_driver(settings) as driver,
        driver.session(database=settings.database) as session,
    ):
        verify_freeze(
            manifest,
            run_metadata(dataset),
            graph_fingerprint(session),
            audit=audit,
            **manifest["parameters"],
        )
        assert (
            graph_version(session) == metadata["index_version"]
        ), "Index version changed during generation."
    assert hashes == {
        str(p.resolve().relative_to(HERE.parents[1])): sha256_file(p) for p in code_files
    }, "Generation code changed during TEST."
    output["by_category"] = {
        category: {
            "questions": len(rows),
            "correct": sum(r["strict"] for r in rows),
            "accuracy": sum(r["strict"] for r in rows) / len(rows),
        }
        for category in sorted({q["category"] for q in selected})
        for rows in [[r for r in output["questions"] if r["category"] == category]]
    }
    (args.output_dir / f"answers_{args.split}.json").write_text(
        json.dumps(output, ensure_ascii=False, indent=2) + "\n"
    )
    fields = [
        "id",
        "category",
        "strict",
        "error",
        "latency_seconds",
        "attempts",
        "first_pass_strict",
    ]
    with (args.output_dir / f"answers_{args.split}.csv").open(
        "w", newline="", encoding="utf-8"
    ) as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        for row in output["questions"]:
            writer.writerow(
                {
                    **{k: row.get(k) for k in fields if k not in {"attempts", "first_pass_strict"}},
                    "attempts": len(row.get("attempts", [])),
                    "first_pass_strict": bool(
                        row.get("attempts") and row["attempts"][0].get("strict")
                    ),
                }
            )
    print(json.dumps(output["summary"]))


if __name__ == "__main__":
    main()