"""Test our extraction loop against 5 convergence PDFs with ground truth.

Runs each PDF through the extraction loop, compares to convergence's gold,
and prints a summary. Also checks if convergence recipes exist for comparison.

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

import json
import sys
import time
from dataclasses import asdict
from functools import partial
from pathlib import Path

# Force unbuffered output
print = partial(print, flush=True)

sys.path.insert(0, str(Path(__file__).parent.parent))

from agentic.config import AgentConfig
from agentic.models import FieldSchema
from agentic.session import ExtractionRunner

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

# ── Test documents ──────────────────────────────────────────────────────────

TESTS = [
    {
        "name": "Fire Inspection (simple)",
        "pdf": CONVERGENCE / "data/pdfs/fire-inspection/1184410.pdf",
        "gt": CONVERGENCE / "data/ground_truth/fire-inspection/1184410.json",
    },
    {
        "name": "Campaign Finance (medium)",
        "pdf": CONVERGENCE / "data/pdfs/campaign-finance/1008174.pdf",
        "gt": CONVERGENCE / "data/ground_truth/campaign-finance/1008174.json",
    },
    {
        "name": "Restaurant Inspection (medium)",
        "pdf": CONVERGENCE / "data/pdfs/restaurant-inspection/1014985.pdf",
        "gt": CONVERGENCE / "data/ground_truth/restaurant-inspection/1014985.json",
    },
    {
        "name": "Building Permit (medium-high)",
        "pdf": CONVERGENCE / "data/pdfs/building-permit/20416188.pdf",
        "gt": CONVERGENCE / "data/ground_truth/building-permit/20416188.json",
    },
    {
        "name": "Arrest Report (complex)",
        "pdf": CONVERGENCE / "data/pdfs/arrest-report/1310759.pdf",
        "gt": CONVERGENCE / "data/ground_truth/arrest-report/1310759.json",
    },
]


def load_convergence_gold(gt_path: Path) -> tuple[list[FieldSchema], dict]:
    """Load convergence ground truth and convert to our schema + gold format."""
    gt = json.loads(gt_path.read_text())
    fields_data = gt.get("fields", {})

    schema = []
    gold = {}
    for name, info in fields_data.items():
        if isinstance(info, dict):
            field_type = info.get("field_type", "string")
            value = info.get("value")
        else:
            field_type = "string"
            value = info

        schema.append(FieldSchema(
            name=name,
            description="",
            field_type=field_type,
        ))
        gold[name] = value

    return schema, gold


def run_test(test: dict, config: AgentConfig) -> dict:
    """Run one test and return results."""
    print(f"\n{'='*70}")
    print(f"  {test['name']}")
    print(f"  PDF: {test['pdf']}")
    print(f"{'='*70}")

    schema, gold = load_convergence_gold(test["gt"])
    print(f"  Fields: {len(schema)}")
    print(f"  Gold keys: {list(gold.keys())[:5]}{'...' if len(gold) > 5 else ''}")

    runner = ExtractionRunner(config)
    session = runner.run(
        pdf_path=str(test["pdf"]),
        page_number=0,
        schema=schema,
        gold=gold,
    )

    # Print results
    print(f"\n  Result: {session.stop_reason}")
    print(f"  Best score: {session.best_score:.0%}")
    print(f"  Iterations: {len(session.attempts)}")
    print(f"  Total tokens: {session.total_tokens}")
    print(f"  Total time: {session.total_time_ms/1000:.1f}s")

    for attempt in session.attempts:
        correct = sum(1 for c in attempt.field_comparisons if c.match)
        total = len(attempt.field_comparisons)
        print(f"    Attempt {attempt.iteration+1}: {correct}/{total} fields ({attempt.overall_score:.0%})")
        if attempt.execution.error_type:
            print(f"      Error: {attempt.execution.error_type} — {attempt.execution.error_message[:100]}")

    # Show best attempt field details
    if session.attempts:
        best = max(session.attempts, key=lambda a: a.overall_score)
        print(f"\n  Best attempt fields:")
        for c in best.field_comparisons:
            status = "✓" if c.match else "✗"
            actual_short = str(c.raw_actual)[:50] if c.raw_actual else "None"
            expected_short = str(c.raw_expected)[:50]
            print(f"    {status} {c.field_name}: {actual_short}")
            if not c.match:
                print(f"      expected: {expected_short}")

    # Save session
    output_path = Path(config.output_dir) / f"convergence_{test['name'].split('(')[0].strip().lower().replace(' ', '_')}.json"
    session.to_json(output_path)

    return {
        "name": test["name"],
        "fields": len(schema),
        "best_score": session.best_score,
        "iterations": len(session.attempts),
        "tokens": session.total_tokens,
        "time_s": session.total_time_ms / 1000,
        "converged": session.converged,
    }


def main():
    config = AgentConfig(
        model="claude-sonnet-4-6",
        max_tokens=16384,
        max_iterations=4,  # 4 for testing speed, can bump later
        target_score=1.0,
        python_command=str(Path(__file__).parent / ".venv" / "bin" / "python"),
        output_dir=str(Path(__file__).parent / "output"),
    )

    results = []
    for test in TESTS:
        if not test["pdf"].exists():
            print(f"\nSKIPPED: {test['name']} — PDF not found: {test['pdf']}")
            continue
        if not test["gt"].exists():
            print(f"\nSKIPPED: {test['name']} — Ground truth not found: {test['gt']}")
            continue

        result = run_test(test, config)
        results.append(result)

    # Summary table
    print(f"\n\n{'='*70}")
    print("  SUMMARY")
    print(f"{'='*70}")
    print(f"  {'Document':<35} {'Fields':>6} {'Score':>6} {'Iters':>5} {'Tokens':>7} {'Time':>6}")
    print(f"  {'-'*35} {'-'*6} {'-'*6} {'-'*5} {'-'*7} {'-'*6}")
    for r in results:
        print(f"  {r['name']:<35} {r['fields']:>6} {r['best_score']:>5.0%} {r['iterations']:>5} {r['tokens']:>7} {r['time_s']:>5.1f}s")


if __name__ == "__main__":
    main()
