"""Compare a v1 and a v2 Azure AI Search index before the alias swap (PB-RAG-01-03 steps 3, 10 and 11).

For each index it reports:
- document count, distinct parent documents, chunks scanned and the largest chunk in characters;
- recall@10 and nDCG@10 on the golden set, scored at parent-document level.
  Chunk keys change when chunking changes, so chunk-level labels cannot compare v1 with v2 (PB-RAG-02-01).

Usage:
    python compare_indexes.py --v1 kb-v1 --v2 kb-v2 --queries eval/queries.json --qrels eval/qrels.json \
        [--vector-field chunk_vector] [--semantic-config default] [--out eval/compare-v2.json]
    python compare_indexes.py --v1 kb-v1 --queries ... --qrels ...      # baseline only (step 3)

queries.json: {"q01": "query text", ...}
qrels.json:   {"q01": {"<parent key>": 2, "<other parent key>": 1}, ...}   (ranx format, graded)

Needs SEARCH_ENDPOINT, a Microsoft Entra ID sign-in (DefaultAzureCredential) with Search Index Data Reader,
and azure-search-documents and azure-identity.
Exit code: 0 pass, 1 v2 fails the gate (parent counts differ or a metric drops more than --tolerance),
2 the input is unreadable.
"""
import argparse
import json
import math
import os
import sys

K = 10
METRICS = ("recall@10", "ndcg@10")
TOLERANCE = 0.02  # team default, same as PB-RAG-02-01; agree it in the change ticket


def parent_ranking(hits, parent_field, k=K):
    """Fold ranked chunk hits into ranked parent keys: first hit of each parent wins."""
    ranked = []
    for hit in hits:
        parent = hit.get(parent_field)
        if parent is not None and parent not in ranked:
            ranked.append(parent)
            if len(ranked) == k:
                break
    return ranked


def recall_at_k(ranked, relevant, k=K):
    wanted = {doc for doc, grade in relevant.items() if grade > 0}
    if not wanted:
        return 0.0
    return len(wanted & set(ranked[:k])) / len(wanted)


def ndcg_at_k(ranked, relevant, k=K):
    """Linear gain, log2 discount (the ranx "ndcg" formula)."""
    dcg = sum(relevant.get(doc, 0) / math.log2(i + 2) for i, doc in enumerate(ranked[:k]))
    ideal = sorted((g for g in relevant.values() if g > 0), reverse=True)[:k]
    idcg = sum(g / math.log2(i + 2) for i, g in enumerate(ideal))
    return dcg / idcg if idcg else 0.0


def search_kwargs(text, parent_field, fetch, vector_field=None, semantic_config=None):
    """Build one query in the production shape: keyword, plus vector (hybrid) and semantic if set."""
    kwargs = {"search_text": text, "top": fetch, "select": [parent_field]}
    if vector_field:
        from azure.search.documents.models import VectorizableTextQuery

        kwargs["vector_queries"] = [
            VectorizableTextQuery(text=text, k_nearest_neighbors=fetch, fields=vector_field)
        ]
    if semantic_config:
        kwargs["query_type"] = "semantic"
        kwargs["semantic_configuration_name"] = semantic_config
    return kwargs


def score(client, queries, qrels, parent_field, fetch=50, vector_field=None, semantic_config=None):
    per_query = {}
    for qid, relevant in qrels.items():
        kwargs = search_kwargs(queries[qid], parent_field, fetch, vector_field, semantic_config)
        ranked = parent_ranking(client.search(**kwargs), parent_field)
        per_query[qid] = {
            "top": ranked,
            "recall@10": recall_at_k(ranked, relevant),
            "ndcg@10": ndcg_at_k(ranked, relevant),
        }
    result = {m: sum(q[m] for q in per_query.values()) / max(len(per_query), 1) for m in METRICS}
    result["per_query"] = per_query
    return result


def index_stats(client, parent_field, chunk_field):
    """Scan every chunk once. Fine for a sandbox; on a large index check the paging limits first."""
    parents, chunks, largest = set(), 0, 0
    for doc in client.search(search_text="*", select=[parent_field, chunk_field]):
        chunks += 1
        if doc.get(parent_field) is not None:
            parents.add(doc[parent_field])
        largest = max(largest, len(doc.get(chunk_field) or ""))
    return {
        "documents": client.get_document_count(),
        "chunks_scanned": chunks,
        "distinct_parents": len(parents),
        "largest_chunk_chars": largest,
    }


def gate(v1, v2, tolerance=TOLERANCE):
    """Return the reasons v2 must not take the alias. Empty list means pass."""
    problems = []
    if v1["distinct_parents"] != v2["distinct_parents"]:
        problems.append(f"distinct parents differ: v1 {v1['distinct_parents']}, v2 {v2['distinct_parents']}")
    for m in METRICS:
        if v2[m] < v1[m] - tolerance:
            problems.append(f"{m} dropped: v1 {v1[m]:.3f}, v2 {v2[m]:.3f}, tolerance {tolerance}")
    return problems


def warnings(report):
    out = []
    for name, side in report["indexes"].items():
        if side["chunks_scanned"] != side["documents"]:
            out.append(f"{name}: scanned {side['chunks_scanned']} of {side['documents']} documents; "
                       "counts are partial (paging limit or writes during the scan)")
        missing = [q for q, r in side["per_query"].items() if not r["top"]]
        if missing:
            out.append(f"{name}: no results for {len(missing)} queries: {', '.join(missing[:5])}")
    return out


def compare(clients, queries, qrels, parent_field="parent_id", chunk_field="chunk",
            fetch=50, vector_field=None, semantic_config=None, tolerance=TOLERANCE):
    """clients: {"v1": client, "v2": client} (v2 optional). Each client has search() and get_document_count()."""
    report = {"tolerance": tolerance, "indexes": {}, "problems": []}
    for name, client in clients.items():
        side = index_stats(client, parent_field, chunk_field)
        side.update(score(client, queries, qrels, parent_field, fetch, vector_field, semantic_config))
        report["indexes"][name] = side
    if "v2" in report["indexes"]:
        report["problems"] = gate(report["indexes"]["v1"], report["indexes"]["v2"], tolerance)
    report["warnings"] = warnings(report)
    return report


def print_table(report, names):
    rows = ("documents", "distinct_parents", "chunks_scanned", "largest_chunk_chars") + METRICS
    sides = report["indexes"]
    print(f"{'':22}" + "".join(f"{n:>16}" for n in names))
    for row in rows:
        cells = [sides[key][row] for key in sides]
        print(f"{row:22}" + "".join(f"{c:>16.3f}" if isinstance(c, float) else f"{c:>16}" for c in cells))
    for w in report["warnings"]:
        print(f"WARNING {w}")
    for p in report["problems"]:
        print(f"FAIL {p}")
    if "v2" in sides:
        print("RESULT: " + ("do not swap" if report["problems"] else "pass"))


def make_client(endpoint, index):
    from azure.identity import DefaultAzureCredential
    from azure.search.documents import SearchClient

    return SearchClient(endpoint, index, DefaultAzureCredential())


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--v1", required=True, help="current index name, for example kb-v1")
    ap.add_argument("--v2", help="new index name; leave out to record the v1 baseline only")
    ap.add_argument("--queries", required=True, help="queries.json: {query id: text}")
    ap.add_argument("--qrels", required=True, help="qrels.json: {query id: {parent key: grade}}")
    ap.add_argument("--parent-field", default="parent_id")
    ap.add_argument("--chunk-field", default="chunk")
    ap.add_argument("--vector-field", help="vector field for a hybrid query (needs a vectorizer on the index)")
    ap.add_argument("--semantic-config", help="semantic configuration name, if production uses semantic ranking")
    ap.add_argument("--fetch", type=int, default=50, help="chunks fetched per query before folding to parents")
    ap.add_argument("--tolerance", type=float, default=TOLERANCE)
    ap.add_argument("--out", help="write the full report as JSON")
    args = ap.parse_args(argv)

    try:
        with open(args.queries, encoding="utf-8") as f:
            queries = json.load(f)
        with open(args.qrels, encoding="utf-8") as f:
            qrels = json.load(f)
        endpoint = os.environ["SEARCH_ENDPOINT"]
    except (OSError, ValueError, KeyError) as e:
        print(f"cannot read input: {e}", file=sys.stderr)
        return 2
    unknown = sorted(set(qrels) - set(queries))
    if unknown:
        print(f"qrels name queries not in {args.queries}: {', '.join(unknown[:5])}", file=sys.stderr)
        return 2

    names = {"v1": args.v1}
    if args.v2:
        names["v2"] = args.v2
    clients = {key: make_client(endpoint, index) for key, index in names.items()}
    report = compare(clients, queries, qrels, args.parent_field, args.chunk_field, args.fetch,
                     args.vector_field, args.semantic_config, args.tolerance)
    report["index_names"] = names
    print_table(report, list(names.values()))
    if args.out:
        with open(args.out, "w", encoding="utf-8") as f:
            json.dump(report, f, indent=2)
    return 1 if report["problems"] else 0


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