"""Test extraction loop on digital PDFs — all in parallel.

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

import asyncio
import json
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from functools import partial
from multiprocessing import Process, Queue
from pathlib import Path

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

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

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


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

    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

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

    pdf_path = str(CONVERGENCE / "data/pdfs" / category / f"{doc_id}.pdf")
    runner = ExtractionRunner(config)
    session = runner.run(pdf_path=pdf_path, schema=schema, gold=gold)
    session.to_json(Path(config.output_dir) / f"digital_{category}.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("Launching", len(TESTS), "extractions in parallel...", flush=True)
    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}", flush=True)

    # Collect results as they finish
    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", flush=True)

    elapsed = time.monotonic() - start

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


if __name__ == "__main__":
    main()
