"""Replay logged app search requests against an index and report where the expected chunks rank (PB-RAG-08-03).

Usage:
    python replay_queries.py requests.jsonl --endpoint https://${SEARCH_SERVICE}.search.windows.net \
        --index ${SANDBOX_INDEX} [--baseline baseline.json] [--out results.json]

requests.jsonl has one logged request per line, exactly as the application sent it:
    {"query_id": "q-001", "body": {"search": "...", "queryType": "semantic", ...},
     "expected_ids": ["doc-12_chunk-3"]}

Authentication uses Microsoft Entra ID (DefaultAzureCredential, role Search Index Data Reader).
The script only queries; it never writes to the index.

Exit code: 0 if every expected chunk is in the top --k, 1 if any is not, 2 if the input is unreadable.
"""
import argparse
import json
import sys

SCOPE = "https://search.azure.com/.default"


def prepare(body, top, debug):
    """Keep the app's request as it is; only widen top and add debug so the rank is visible."""
    body = dict(body)
    body["top"] = max(int(body.get("top", 0)), top)
    if debug:
        body["debug"] = debug
    return body


def rank_report(results, expected_ids, key):
    """Return {expected_id: {"rank": 1-based or None, scores...}} for one result list."""
    report = {}
    for doc_id in expected_ids:
        report[doc_id] = {"rank": None}
        for rank, doc in enumerate(results, start=1):
            if doc.get(key) == doc_id:
                subscores = doc.get("@search.documentDebugInfo", {}).get("vectors", {}).get("subscores", {})
                report[doc_id] = {
                    "rank": rank,
                    "score": doc.get("@search.score"),
                    "reranker_score": doc.get("@search.rerankerScore"),
                    "text_score": subscores.get("text", {}).get("searchScore"),
                    "vector_scores": subscores.get("vectors"),
                }
                break
    return report


def compare(current, baseline, k):
    """List (query_id, doc_id, old_rank, new_rank) where a chunk fell out of the top k."""
    drops = []
    for qid, docs in current.items():
        for doc_id, now in docs.items():
            old = baseline.get(qid, {}).get(doc_id, {}).get("rank")
            new = now["rank"]
            if (new is None or new > k) and (old is not None and old <= k):
                drops.append((qid, doc_id, old, new))
    return drops


def search(endpoint, index, api_version, body, token):
    import requests  # imported here so the pure functions above need no packages

    url = f"{endpoint}/indexes/{index}/docs/search?api-version={api_version}"
    resp = requests.post(url, json=body, headers={"Authorization": f"Bearer {token}"}, timeout=30)
    resp.raise_for_status()
    return resp.json().get("value", [])


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("requests_file")
    p.add_argument("--endpoint", required=True)
    p.add_argument("--index", required=True, help="Sandbox copy, or an alias; never tune on production")
    p.add_argument("--api-version", default="2026-04-01")
    p.add_argument("--key", default="id", help="Key field of the index")
    p.add_argument("--top", type=int, default=50)
    p.add_argument("--k", type=int, default=10, help="Cut-off that counts as 'found'")
    p.add_argument("--debug", default="all", help="debug mode; use '' to send none")
    p.add_argument("--baseline", help="results.json from the last good run")
    p.add_argument("--out", default="results.json")
    args = p.parse_args(argv)

    try:
        with open(args.requests_file, encoding="utf-8") as f:
            rows = [json.loads(line) for line in f if line.strip()]
        baseline = json.load(open(args.baseline, encoding="utf-8")) if args.baseline else {}
    except (OSError, ValueError) as err:
        print(f"Cannot read input: {err}", file=sys.stderr)
        return 2

    from azure.identity import DefaultAzureCredential

    token = DefaultAzureCredential().get_token(SCOPE).token
    current = {}
    for row in rows:
        body = prepare(row["body"], args.top, args.debug)
        results = search(args.endpoint, args.index, args.api_version, body, token)
        current[row["query_id"]] = rank_report(results, row["expected_ids"], args.key)

    with open(args.out, "w", encoding="utf-8") as f:
        json.dump(current, f, indent=2)

    missing = [(q, d) for q, docs in current.items() for d, r in docs.items() if r["rank"] is None or r["rank"] > args.k]
    print(f"{len(rows)} queries, {len(missing)} expected chunks outside the top {args.k}")
    for q, d in missing:
        print(f"  {q}: {d} rank={current[q][d]['rank']}")
    for q, d, old, new in compare(current, baseline, args.k):
        print(f"  DROPPED {q}: {d} rank {old} -> {new}")
    return 1 if missing else 0


if __name__ == "__main__":
    sys.exit(main())