"""Unified test suite for agentic extraction.

Runs all configured PDFs through the extraction loop, supports parallel execution,
cached gold, and produces a summary table.

Run: agentic/.venv/bin/python agentic/test_suite.py
     agentic/.venv/bin/python agentic/test_suite.py --filter autopsy
     agentic/.venv/bin/python agentic/test_suite.py --max-iters 4 --sequential
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from functools import partial
from pathlib import Path

print = partial(print, flush=True)
sys.path.insert(0, str(Path(__file__).parent.parent))

CONVERGENCE = Path("/Users/soma/Development/convergence")
PRACTICE_PDF = Path(__file__).parent.parent / "pdfs" / "01-practice.pdf"

# ── Test registry ───────────────────────────────────────────────────────────
# Each entry: (source, category, doc_id, gold_source)
#   source: "convergence" (uses CONVERGENCE path) or "local" (uses project pdfs/)
#   gold_source: "ground_truth" | "llm_cached" | "llm_generate" | "manual"

TESTS = [
    # Round 1 — convergence PDFs with ground truth
    ("convergence", "restaurant-inspection", "1155156", "ground_truth"),
    ("convergence", "traffic-accident", "1195902", "ground_truth"),
    ("convergence", "meeting-minutes", "1200271", "ground_truth"),
    ("convergence", "financial-disclosure", "1194193", "ground_truth"),
    # Round 2 — convergence PDFs with LLM-generated gold (cached)
    ("convergence", "campaign-finance", "23117530", "llm_cached"),
    ("convergence", "autopsy-report", "6543431", "llm_cached"),
    ("convergence", "search-warrant", "21062872", "llm_cached"),
    ("convergence", "election-result", "25172686", "llm_cached"),
    # Round 3 — new categories
    ("convergence", "use-of-force", "21099220", "ground_truth"),
    ("convergence", "environmental-permit", "21060199", "llm_generate"),
    ("convergence", "death-certificate", "22130594", "llm_generate"),
    ("convergence", "lobbying-disclosure", "1003922", "llm_generate"),
    ("convergence", "business-license", "685794", "llm_generate"),
]


# ── Gold loading ────────────────────────────────────────────────────────────

def load_gold_from_ground_truth(category: str, doc_id: str) -> tuple[list[dict], dict]:
    """Load gold from convergence ground truth JSON."""
    gt_path = CONVERGENCE / "data/ground_truth" / category / f"{doc_id}.json"
    gt = json.loads(gt_path.read_text())
    schema = []
    gold = {}
    for name, info in gt.get("fields", {}).items():
        ft = info.get("field_type", "string") if isinstance(info, dict) else "string"
        val = info.get("value") if isinstance(info, dict) else info
        schema.append({"name": name, "field_type": ft})
        gold[name] = val
    return schema, gold


def load_gold_from_cache(category: str, doc_id: str) -> tuple[list[dict], dict]:
    """Load gold from cached LLM-generated gold."""
    cache_path = Path(__file__).parent / "output" / "gold" / category / f"{doc_id}.json"
    if not cache_path.exists():
        raise FileNotFoundError(f"No cached gold at {cache_path}. Run with --generate-gold first.")
    data = json.loads(cache_path.read_text())
    return data["schema"], data["gold"]


def load_gold(category: str, doc_id: str, gold_source: str, pdf_path: str) -> tuple[list[dict], dict]:
    """Load gold from the appropriate source."""
    if gold_source == "ground_truth":
        return load_gold_from_ground_truth(category, doc_id)
    elif gold_source in ("llm_cached", "llm_generate"):
        try:
            return load_gold_from_cache(category, doc_id)
        except FileNotFoundError:
            if gold_source == "llm_generate":
                from agentic.config import AgentConfig
                from agentic.gold import generate_gold_via_llm
                from agentic.models import FieldSchema
                schema_objs, gold_obj = generate_gold_via_llm(pdf_path, AgentConfig())
                schema = [{"name": f.name, "field_type": f.field_type} for f in schema_objs]
                return schema, gold_obj.fields
            raise
    else:
        raise ValueError(f"Unknown gold_source: {gold_source}")


# ── Worker ──────────────────────────────────────────────────────────────────

def run_one(
    source: str,
    category: str,
    doc_id: str,
    gold_source: str,
    max_iters: int,
) -> dict:
    """Run extraction for one document. Returns result dict."""
    from agentic.config import AgentConfig
    from agentic.models import FieldSchema
    from agentic.session import ExtractionRunner

    doc_name = f"{category}/{doc_id}"

    if source == "convergence":
        pdf_path = str(CONVERGENCE / "data/pdfs" / category / f"{doc_id}.pdf")
    else:
        pdf_path = str(PRACTICE_PDF)

    try:
        schema_dicts, gold = load_gold(category, doc_id, gold_source, pdf_path)
    except Exception as e:
        return {
            "doc": doc_name,
            "fields": 0,
            "best": 0.0,
            "iters": 0,
            "tokens": 0,
            "time": 0,
            "attempts": [],
            "error": str(e),
        }

    schema = [FieldSchema(**f) for f in schema_dicts]

    config = AgentConfig(
        model="claude-sonnet-4-6",
        max_tokens=16384,
        max_iterations=max_iters,
        python_command=str(Path(__file__).parent / ".venv" / "bin" / "python"),
        output_dir=str(Path(__file__).parent / "output"),
    )

    runner = ExtractionRunner(config)
    session = runner.run(pdf_path=pdf_path, schema=schema, gold=gold)
    session.to_json(Path(config.output_dir) / f"suite_{category}_{doc_id}.json")

    return {
        "doc": doc_name,
        "fields": len(schema),
        "best": session.best_score,
        "iters": len(session.attempts),
        "tokens": session.total_tokens,
        "time": session.total_time_ms / 1000,
        "converged": session.converged,
        "stop_reason": session.stop_reason,
        "attempts": [
            {"iter": a.iteration + 1, "score": a.overall_score, "approach": a.diagnostics.approach_used}
            for a in session.attempts
        ],
    }


# ── Runner ──────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Agentic extraction test suite")
    parser.add_argument("--filter", type=str, help="Only run tests matching this substring")
    parser.add_argument("--max-iters", type=int, default=4, help="Max iterations per document")
    parser.add_argument("--sequential", action="store_true", help="Run sequentially instead of parallel")
    parser.add_argument("--workers", type=int, default=4, help="Max parallel workers (default 4)")
    parser.add_argument("--generate-gold", action="store_true", help="Generate gold via LLM if not cached")
    args = parser.parse_args()

    # Filter tests
    tests = TESTS
    if args.filter:
        tests = [(s, c, d, g) for s, c, d, g in tests if args.filter.lower() in f"{c}/{d}".lower()]

    if not tests:
        print("No tests match filter.")
        return

    # Upgrade llm_cached to llm_generate if flag is set
    if args.generate_gold:
        tests = [(s, c, d, "llm_generate" if g == "llm_cached" else g) for s, c, d, g in tests]

    workers = 1 if args.sequential else args.workers
    print(f"Running {len(tests)} extraction tests (max {args.max_iters} iterations, {workers} workers)")
    print()

    start = time.monotonic()
    results = []

    if args.sequential:
        for source, category, doc_id, gold_source in tests:
            print(f"  Running: {category}/{doc_id}...")
            result = run_one(source, category, doc_id, gold_source, args.max_iters)
            results.append(result)
            trajectory = " -> ".join(f"{a['score']:.0%}" for a in result["attempts"])
            print(f"  Done: {result['doc']} — {result['best']:.0%} ({trajectory}) in {result['time']:.0f}s")
    else:
        # Bounded concurrency via ProcessPoolExecutor
        with ProcessPoolExecutor(max_workers=workers) as executor:
            futures = {}
            for source, category, doc_id, gold_source in tests:
                future = executor.submit(
                    run_one, source, category, doc_id, gold_source, args.max_iters
                )
                futures[future] = f"{category}/{doc_id}"
                print(f"  Queued: {category}/{doc_id}")

            for future in as_completed(futures):
                name = futures[future]
                try:
                    result = future.result()
                except Exception as e:
                    result = {"doc": name, "fields": 0, "best": 0.0, "iters": 0,
                              "tokens": 0, "time": 0, "attempts": [], "error": str(e)}
                results.append(result)
                trajectory = " -> ".join(f"{a['score']:.0%}" for a in result["attempts"])
                status = "ERROR" if result.get("error") else f"{result['best']:.0%}"
                print(f"  Done: {result['doc']} — {status} ({trajectory}) in {result['time']:.0f}s")

    elapsed = time.monotonic() - start

    # ── Summary ─────────────────────────────────────────────────────────────
    print(f"\n{'='*72}")
    print(f"  RESULTS ({len(results)} docs, {elapsed:.0f}s wall time)")
    print(f"{'='*72}")
    print(f"  {'Document':<35} {'Fields':>6} {'Score':>6} {'Iters':>5} {'Time':>6} {'Status':<12}")
    print(f"  {'-'*35} {'-'*6} {'-'*6} {'-'*5} {'-'*6} {'-'*12}")

    for r in sorted(results, key=lambda r: -r["best"]):
        error = r.get("error")
        if error:
            status = f"ERROR: {error[:30]}"
        elif r.get("converged"):
            status = "converged"
        else:
            status = r.get("stop_reason", "")
        print(f"  {r['doc']:<35} {r['fields']:>6} {r['best']:>5.0%} {r['iters']:>5} {r['time']:>5.0f}s {status:<12}")

    # Stats
    valid = [r for r in results if not r.get("error")]
    if valid:
        avg = sum(r["best"] for r in valid) / len(valid)
        perfect = sum(1 for r in valid if r["best"] >= 1.0)
        total_tokens = sum(r["tokens"] for r in valid)
        print(f"\n  Average: {avg:.0%} | Perfect: {perfect}/{len(valid)} | Tokens: {total_tokens:,}")

    # Save summary
    summary_path = Path(__file__).parent / "output" / "suite_summary.json"
    summary_path.write_text(json.dumps(results, indent=2, default=str))
    print(f"  Summary saved to {summary_path}")


if __name__ == "__main__":
    main()
