"""Test round 2: 5 more digital PDFs. Generate gold via LLM for those without it.

Run: agentic/.venv/bin/python agentic/test_digital_round2.py
"""

import json
import sys
import time
from functools import partial
from multiprocessing import Process, Queue
from pathlib import Path

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

CONVERGENCE = Path("/Users/soma/Development/convergence")

TESTS = [
    ("campaign-finance", "23117530"),
    ("autopsy-report", "6543431"),
    ("search-warrant", "21062872"),
    ("election-result", "25172686"),
    ("government-contract", "20617970"),
]


def generate_gold_via_llm(pdf_path: str) -> tuple[list[dict], dict]:
    """Use Claude to generate gold standard from the PDF text."""
    import anthropic
    from dotenv import load_dotenv
    load_dotenv(Path(__file__).parent / ".env")

    from natural_pdf import PDF
    pdf = PDF(pdf_path)
    # Get text from all pages
    all_text = "\n\n".join(page.extract_text(layout=True) for page in pdf.pages)
    # Truncate if huge
    if len(all_text) > 15000:
        all_text = all_text[:15000] + "\n... (truncated)"
    pdf.close()

    client = anthropic.Anthropic()
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=4096,
        temperature=0,
        messages=[{
            "role": "user",
            "content": f"""Extract all meaningful fields from this document as JSON.

Return a JSON object where keys are snake_case field names and values are the extracted text.
Include a "field_types" key that maps each field name to its type: "string", "number", "currency", "date", "boolean", or "address".
Only extract fields that have actual values (not blank/empty).
Use the exact text from the document for values — don't paraphrase.

Document text:
{all_text}

Return ONLY the JSON object, no explanation."""
        }],
    )

    raw = response.content[0].text.strip()
    # Strip markdown fences if present
    if raw.startswith("```"):
        raw = "\n".join(raw.split("\n")[1:])
    if raw.endswith("```"):
        raw = raw[:-3]

    data = json.loads(raw.strip())

    # Separate field_types from values
    field_types = data.pop("field_types", {})

    schema = []
    gold = {}
    for name, value in data.items():
        ft = field_types.get(name, "string")
        schema.append({"name": name, "field_type": ft})
        gold[name] = value

    return schema, gold


def run_one(category: str, doc_id: str, queue: Queue):
    """Run extraction for one document."""
    from agentic.config import AgentConfig
    from agentic.models import FieldSchema
    from agentic.session import ExtractionRunner

    pdf_path = str(CONVERGENCE / "data/pdfs" / category / f"{doc_id}.pdf")
    gt_path = CONVERGENCE / "data/ground_truth" / category / f"{doc_id}.json"

    # Load or generate gold
    if gt_path.exists():
        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
        print(f"  {category}/{doc_id}: loaded {len(schema)} fields from existing gold")
    else:
        print(f"  {category}/{doc_id}: generating gold via LLM...")
        schema, gold = generate_gold_via_llm(pdf_path)
        print(f"  {category}/{doc_id}: generated {len(schema)} fields")
        # Save for reuse
        gold_dir = Path(f"agentic/output/gold/{category}")
        gold_dir.mkdir(parents=True, exist_ok=True)
        (gold_dir / f"{doc_id}.json").write_text(json.dumps({"schema": schema, "gold": gold}, indent=2))

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

    config = AgentConfig(
        model="claude-sonnet-4-6",
        max_tokens=16384,
        max_iterations=4,
        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_objs, gold=gold)
    session.to_json(Path(config.output_dir) / f"r2_{category}_{doc_id}.json")

    result = {
        "doc": f"{category}/{doc_id}",
        "fields": len(schema),
        "best": session.best_score,
        "iters": len(session.attempts),
        "tokens": session.total_tokens,
        "time": session.total_time_ms / 1000,
        "attempts": [
            {"iter": a.iteration + 1, "score": a.overall_score, "error": a.execution.error_type}
            for a in session.attempts
        ],
    }
    queue.put(result)


def main():
    print(f"Launching {len(TESTS)} extractions in parallel...")
    start = time.monotonic()

    queue = Queue()
    processes = []
    for category, doc_id in TESTS:
        p = Process(target=run_one, args=(category, doc_id, queue))
        p.start()
        processes.append((p, f"{category}/{doc_id}"))
        print(f"  Started: {category}/{doc_id}")

    results = []
    for p, name in processes:
        p.join()
        if not queue.empty():
            result = queue.get()
            results.append(result)
            attempts = " → ".join(f"{a['score']:.0%}" for a in result["attempts"])
            print(f"  Done: {result['doc']} — {result['best']:.0%} ({attempts}) in {result['time']:.0f}s")

    elapsed = time.monotonic() - start

    print(f"\n{'='*60}")
    print(f"  SUMMARY (wall time: {elapsed:.0f}s)")
    print(f"{'='*60}")
    print(f"  {'Document':<40} {'Fields':>6} {'Score':>6} {'Iters':>5} {'Time':>6}")
    print(f"  {'-'*40} {'-'*6} {'-'*6} {'-'*5} {'-'*6}")
    for r in sorted(results, key=lambda r: -r["best"]):
        print(f"  {r['doc']:<40} {r['fields']:>6} {r['best']:>5.0%} {r['iters']:>5} {r['time']:>5.0f}s")


if __name__ == "__main__":
    main()
