"""Paired compression accuracy and real cold/warm cost measurements.

python -m app.benchmark compression --input data/normalized.jsonl --limit 20
python -m app.benchmark search --idea 'Does vitamin D reduce depression?'
"""

import argparse
import asyncio
import json
from pathlib import Path
from time import perf_counter

import httpx

from app.config import Settings, settings
from app.ingest.__main__ import read_jsonl
from app.llm import LLMService, Usage

FIELDS = ("n", "estimate", "ci_low", "ci_high", "p_value", "effect_type", "outcome")


async def compression_benchmark(path: Path, limit: int) -> dict:
    config = Settings(**settings.model_dump())
    config.compression_enabled = True
    config.compression_validated = True  # Experiment only; never changes live configuration.
    if not config.ttc_api_key or not config.openai_api_key:
        raise ValueError("OpenAI and TTC keys are required for the paired benchmark")
    llm = LLMService(config)
    rows = []
    try:
        for study in read_jsonl(path):
            if len(rows) >= limit:
                break
            if (
                not study.get("abstract")
                or study.get("is_review")
                or study.get("source") != "openalex"
            ):
                continue
            pair = {"id": study["id"]}
            for compressed in (False, True):
                usage = Usage()
                label = "compressed" if compressed else "original"
                started = perf_counter()
                try:
                    result = await llm.extract(study, usage, compression=compressed)
                    pair[label] = {
                        "verified": True,
                        "facts": {key: result.get(key) for key in FIELDS},
                    }
                except Exception as exc:
                    pair[label] = {"verified": False, "failure": type(exc).__name__, "facts": {}}
                pair[label]["usage"] = usage.summary([study["abstract"]], config)
                pair[label]["wallMs"] = round((perf_counter() - started) * 1000)
            pair["exactFactAgreement"] = pair["original"]["facts"] == pair["compressed"]["facts"]
            rows.append(pair)
    finally:
        await llm.close()
    passed = sum(
        r["original"]["verified"] and r["compressed"]["verified"] and r["exactFactAgreement"]
        for r in rows
    )
    return {
        "kind": "paired_compression",
        "studies": len(rows),
        "verifiedExactAgreements": passed,
        "agreementRate": passed / len(rows) if rows else None,
        "enableRecommended": bool(rows) and passed == len(rows) and config.ttc_input_price > 0,
        "limitation": "Agreement with uncompressed model extraction is not human-validated clinical accuracy. Review every difference and include actual TTC pricing before enabling compression.",
        "rows": rows,
    }


async def search_benchmark(idea: str, url: str) -> dict:
    rows = []
    async with httpx.AsyncClient(timeout=300) as client:
        for label in ("first", "repeat"):
            started = perf_counter()
            response = await client.post(
                f"{url.rstrip('/')}/search", json={"idea": idea, "sesoi": 0.2}
            )
            response.raise_for_status()
            result = response.json()
            rows.append(
                {
                    "run": label,
                    "wallMs": round((perf_counter() - started) * 1000),
                    "matched": result["totalScanned"],
                    "retrieval": result["retrieval"],
                    "costs": result["costs"],
                    "warnings": result["warnings"],
                }
            )
    return {
        "kind": "measured_repeated_search",
        "idea": idea,
        "runs": rows,
        "note": "The first run is cold only for documents not previously extracted. Inspect cache hits.",
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)
    compression = sub.add_parser("compression")
    compression.add_argument("--input", type=Path, required=True)
    compression.add_argument("--limit", type=int, default=20)
    search = sub.add_parser("search")
    search.add_argument("--idea", required=True)
    search.add_argument("--url", default="http://127.0.0.1:8000")
    for command in (compression, search):
        command.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if getattr(args, "limit", 1) <= 0:
        parser.error("--limit must be positive")
    result = asyncio.run(
        compression_benchmark(args.input, args.limit)
        if args.command == "compression"
        else search_benchmark(args.idea, args.url)
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2))
    print(f"Benchmark written to {args.output}")


if __name__ == "__main__":
    main()