"""Compare the two playground rankers on one document's questions.

For a document key and a question list, runs `retrieve()` once per ranker,
has the playground answer agent answer from each context, and writes one JSON
row per question: both chunk lists with their fused positions, the Jev page
verdicts and chunk scores, both answers with their cited ids, timings and
failure counts. A measurement tool, not part of the package.

    uv run python experiments/jev_ranking/compare.py <doc_key> <questions.json> <out.json> [--only jev|haiku]

`questions.json` is a JSON list of question strings, or `sweep` to take the
document's cached sweep questions from `ade_playground.sweep_answers`.
"""

from __future__ import annotations

import asyncio
import json
import sys
import time
from pathlib import Path
from typing import Any, Dict, List

from dotenv import load_dotenv

load_dotenv()

from quber.playground import db  # noqa: E402
from quber.playground.agent import answer, build_context  # noqa: E402
from quber.playground.answers.document import DocumentIdentity  # noqa: E402
from quber.playground.ranking import score  # noqa: E402
from quber.playground.retrieval import FUSED_LIMIT, RetrievedChunk, _fused_window, retrieve  # noqa: E402
from quber.settings import get_settings  # noqa: E402

K = 10


def identity(doc_key: str) -> DocumentIdentity:
    with db.connect() as conn:
        r = conn.execute(
            "SELECT title, folder, filing_type, period, year FROM ade_playground.documents WHERE doc_key=%s",
            (doc_key,),
        ).fetchone()
    return DocumentIdentity(title=r[0], issuer=r[1], form=r[2], period=r[3], year=r[4])


def sweep_questions(doc_key: str) -> List[str]:
    with db.connect() as conn:
        rows = conn.execute(
            "SELECT question FROM ade_playground.sweep_answers WHERE doc_key=%s AND valid_to IS NULL ORDER BY id",
            (doc_key,),
        ).fetchall()
    return [r[0] for r in rows]


async def arm(
    ranker: str, doc_key: str, question: str, ident: DocumentIdentity, fused: List[RetrievedChunk]
) -> Dict:
    get_settings().playground.ranker = ranker  # type: ignore[misc]
    position = {c.chunk_id: n + 1 for n, c in enumerate(fused)}
    t0 = time.perf_counter()
    chunks = await retrieve(doc_key, question, k=K)
    t1 = time.perf_counter()
    try:
        a = await answer(question, chunks, ident)
        reply: Dict[str, Any] = {"answer": a.answer, "cited_ids": a.cited_ids}
    except Exception as exc:  # recorded, not hidden
        reply = {"error": repr(exc)}
    return {
        "chunk_ids": [c.chunk_id for c in chunks],
        "pages": [c.page + 1 for c in chunks],
        "fused_positions": [position.get(c.chunk_id) for c in chunks],
        "context_chars": len(build_context(chunks)),
        "retrieve_seconds": round(t1 - t0, 2),
        "answer_seconds": round(time.perf_counter() - t1, 2),
        **reply,
    }


async def jev_detail(question: str, fused: List[RetrievedChunk]) -> Dict:
    r = await score(question, fused)
    return {
        "pages": [
            {"page": p.page + 1, "relevant": p.relevant, "confidence": p.confidence, "error": p.error}
            for p in r.pages
        ],
        "scores": [
            {
                "chunk_id": s.chunk_id,
                "page": s.page + 1,
                "fused_position": s.index + 1,
                "grade": s.grade,
                "score": s.score,
                "error": s.error,
            }
            for s in sorted(r.scores, key=lambda s: -(s.score or -1))
        ],
        "seconds": round(r.seconds, 2),
        "input_tokens": r.input_tokens,
        "failures": r.failures,
        "model": r.model,
    }


async def main() -> None:
    doc_key, qfile, out = sys.argv[1], sys.argv[2], Path(sys.argv[3])
    only = sys.argv[sys.argv.index("--only") + 1] if "--only" in sys.argv else None
    questions = sweep_questions(doc_key) if qfile == "sweep" else json.load(open(qfile))
    ident = identity(doc_key)
    rows = []
    for n, q in enumerate(questions):
        fused = await _fused_window(doc_key, q, FUSED_LIMIT)
        row: Dict[str, Any] = {"question": q, "fused_count": len(fused)}
        if only != "haiku":
            row["jev"] = await arm("jev", doc_key, q, ident, fused)
            row["jev_detail"] = await jev_detail(q, fused)
        if only != "jev":
            row["haiku"] = await arm("haiku", doc_key, q, ident, fused)
        rows.append(row)
        out.write_text(json.dumps(rows, indent=1))
        print(
            f"[{n + 1}/{len(questions)}] "
            + " ".join(
                f"{k}={v['retrieve_seconds']}s"
                for k, v in row.items()
                if isinstance(v, dict) and "retrieve_seconds" in v
            ),
            flush=True,
        )


asyncio.run(main())