"""Test extraction loop on digital PDFs from convergence.

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

import json
import sys
from functools import partial
from pathlib import Path

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")

TESTS = [
    ("financial-disclosure", "1194193"),
    ("restaurant-inspection", "1155156"),
    ("search-warrant", "21092503"),
    ("traffic-accident", "1195902"),
    ("meeting-minutes", "1200271"),
]

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"),
)


def load_gold(category, doc_id):
    gt = json.loads((CONVERGENCE / "data/ground_truth" / category / f"{doc_id}.json").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(FieldSchema(name=name, field_type=ft))
        gold[name] = val
    return schema, gold


def main():
    runner = ExtractionRunner(config)
    results = []

    for category, doc_id in TESTS:
        pdf_path = CONVERGENCE / "data/pdfs" / category / f"{doc_id}.pdf"
        print(f"\n{'='*60}")
        print(f"  {category}/{doc_id}")
        print(f"{'='*60}")

        schema, gold = load_gold(category, doc_id)
        print(f"  Fields: {len(schema)}")

        session = runner.run(
            pdf_path=str(pdf_path),
            schema=schema,
            gold=gold,
        )

        for a in session.attempts:
            correct = sum(1 for c in a.field_comparisons if c.match)
            err = f" Error: {a.execution.error_type}" if a.execution.error_type else ""
            print(f"  Attempt {a.iteration+1}: {correct}/{len(schema)} ({a.overall_score:.0%}){err}")

        results.append({
            "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,
        })

        session.to_json(Path(config.output_dir) / f"digital_{category}.json")

    print(f"\n\n{'='*60}")
    print(f"  SUMMARY")
    print(f"{'='*60}")
    print(f"  {'Document':<35} {'Fields':>6} {'Score':>6} {'Iters':>5} {'Time':>6}")
    print(f"  {'-'*35} {'-'*6} {'-'*6} {'-'*5} {'-'*6}")
    for r in results:
        print(f"  {r['doc']:<35} {r['fields']:>6} {r['best']:>5.0%} {r['iters']:>5} {r['time']:>5.1f}s")


if __name__ == "__main__":
    main()
