#!/usr/bin/env python3
"""
Compare coarse-ranking score components between two indices for queries that regressed
in evaluation reports.

This script answers a narrower question than field diffing:
for the documents that matter in worse queries, did the ranking move because of
image KNN, text KNN, lexical/text score, or coarse-window recall?

Typical usage:
  ./.venv/bin/python scripts/inspect/analyze_coarse_component_regression.py \
    --current-report artifacts/search_evaluation/batch_reports/batch_20260417T073901Z_00b6a8aa3d.json \
    --backup-report artifacts/search_evaluation/batch_reports/batch_20260417T074717Z_00b6a8aa3d.json \
    --current-index search_products_tenant_163 \
    --backup-index search_products_tenant_163_backup_20260415_1438
"""

from __future__ import annotations

import argparse
import logging
import os
import statistics
import sys
from collections import Counter
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence, Tuple

PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from config import get_app_config
from context.request_context import create_request_context
from query import QueryParser
from search import Searcher
from utils.es_client import get_es_client_from_env

from scripts.inspect.analyze_eval_index_regression import _load_report


logger = logging.getLogger("coarse_component_regression")


def _rank_map(rows: Sequence[Dict[str, Any]]) -> Dict[str, int]:
    return {str(row["spu_id"]): int(row["rank"]) for row in rows}


def _collect_regressed_docs(
    current_report: Dict[str, Any],
    backup_report: Dict[str, Any],
    *,
    rank_gap_threshold: int,
    scan_depth: int,
) -> Dict[str, List[Dict[str, Any]]]:
    current_per_query = {row["query"]: row for row in current_report["per_query"]}
    backup_per_query = {row["query"]: row for row in backup_report["per_query"]}
    grouped: Dict[str, List[Dict[str, Any]]] = {}
    for query, current_case in current_per_query.items():
        backup_case = backup_per_query[query]
        delta = (
            float(current_case["metrics"]["Primary_Metric_Score"])
            - float(backup_case["metrics"]["Primary_Metric_Score"])
        )
        if delta >= 0:
            continue
        current_ranks = _rank_map(current_case["top_results"])
        for row in backup_case["top_results"][:scan_depth]:
            if row["label"] not in {"Fully Relevant", "Mostly Relevant"}:
                continue
            current_rank = current_ranks.get(row["spu_id"], 999)
            if current_rank <= int(row["rank"]) + rank_gap_threshold:
                continue
            grouped.setdefault(query, []).append(
                {
                    "query": query,
                    "delta_primary": delta,
                    "spu_id": str(row["spu_id"]),
                    "backup_rank_eval": int(row["rank"]),
                    "backup_label": str(row["label"]),
                    "current_rank_eval": current_rank,
                }
            )
    return grouped


def _build_searcher() -> Searcher:
    config = get_app_config().search
    es_client = get_es_client_from_env()
    query_parser = QueryParser(config)
    return Searcher(es_client, config, query_parser)


def _run_query(searcher: Searcher, *, query: str, tenant_id: str, index_name: str) -> Tuple[Dict[str, Dict[str, Any]], int]:
    os.environ[f"ES_INDEX_OVERRIDE_TENANT_{tenant_id}"] = index_name
    ctx = create_request_context(reqid="coarsecmp", uid="-1")
    ctx._logger = logger
    searcher.search(
        query=query,
        tenant_id=tenant_id,
        size=10,
        context=ctx,
        debug=True,
        enable_rerank=False,
        language="en",
    )
    rows = ctx.get_intermediate_result("coarse_rank_scores", []) or []
    by_doc: Dict[str, Dict[str, Any]] = {}
    for rank, row in enumerate(rows, start=1):
        doc_id = row.get("doc_id")
        if doc_id is None:
            continue
        payload = dict(row)
        payload["_coarse_rank"] = rank
        by_doc[str(doc_id)] = payload
    return by_doc, len(rows)


def _safe_float(value: Any) -> float | None:
    try:
        if value is None:
            return None
        return float(value)
    except (TypeError, ValueError):
        return None


def _delta(current_value: Any, backup_value: Any) -> float | None:
    current = _safe_float(current_value)
    backup = _safe_float(backup_value)
    if current is None or backup is None:
        return None
    return current - backup


def _counter_key(delta_value: float | None, *, eps: float = 1e-6) -> str:
    if delta_value is None:
        return "missing"
    if abs(delta_value) <= eps:
        return "same"
    return "lower" if delta_value < 0 else "higher"


def _median_or_none(values: Sequence[float]) -> float | None:
    if not values:
        return None
    return float(statistics.median(values))


def _summarize_rows(comparisons: Sequence[Dict[str, Any]]) -> None:
    both_present = [row for row in comparisons if row["current_row"] is not None and row["backup_row"] is not None]
    backup_only = [row for row in comparisons if row["current_row"] is None and row["backup_row"] is not None]
    current_only = [row for row in comparisons if row["current_row"] is not None and row["backup_row"] is None]

    image_counter: Counter[str] = Counter()
    text_knn_counter: Counter[str] = Counter()
    text_counter: Counter[str] = Counter()
    es_counter: Counter[str] = Counter()
    coarse_counter: Counter[str] = Counter()

    image_deltas: List[float] = []
    text_knn_deltas: List[float] = []
    text_deltas: List[float] = []
    es_deltas: List[float] = []
    coarse_deltas: List[float] = []

    for row in both_present:
        image_delta = _delta(row["current_row"].get("image_knn_score"), row["backup_row"].get("image_knn_score"))
        text_knn_delta = _delta(row["current_row"].get("text_knn_score"), row["backup_row"].get("text_knn_score"))
        text_delta = _delta(row["current_row"].get("text_score"), row["backup_row"].get("text_score"))
        es_delta = _delta(row["current_row"].get("es_score"), row["backup_row"].get("es_score"))
        coarse_delta = _delta(row["current_row"].get("coarse_score"), row["backup_row"].get("coarse_score"))

        image_counter[_counter_key(image_delta)] += 1
        text_knn_counter[_counter_key(text_knn_delta)] += 1
        text_counter[_counter_key(text_delta)] += 1
        es_counter[_counter_key(es_delta)] += 1
        coarse_counter[_counter_key(coarse_delta)] += 1

        for bucket, sink in (
            (image_delta, image_deltas),
            (text_knn_delta, text_knn_deltas),
            (text_delta, text_deltas),
            (es_delta, es_deltas),
            (coarse_delta, coarse_deltas),
        ):
            if bucket is not None:
                sink.append(bucket)

    print("Coarse Component Summary")
    print("=" * 80)
    print(f"affected_docs: {len(comparisons)}")
    print(f"present_in_both_coarse_windows: {len(both_present)}")
    print(f"only_in_backup_coarse_window: {len(backup_only)}")
    print(f"only_in_current_coarse_window: {len(current_only)}")
    print()
    print(f"image_knn delta buckets: {dict(image_counter)}")
    print(f"text_knn delta buckets : {dict(text_knn_counter)}")
    print(f"text_score delta buckets: {dict(text_counter)}")
    print(f"es_score delta buckets  : {dict(es_counter)}")
    print(f"coarse_score buckets    : {dict(coarse_counter)}")
    print()
    print(
        "median deltas (current - backup): "
        f"image_knn={_median_or_none(image_deltas)} | "
        f"text_knn={_median_or_none(text_knn_deltas)} | "
        f"text_score={_median_or_none(text_deltas)} | "
        f"es_score={_median_or_none(es_deltas)} | "
        f"coarse_score={_median_or_none(coarse_deltas)}"
    )
    print()


def _print_query_examples(comparisons: Sequence[Dict[str, Any]], top_queries: int, docs_per_query: int) -> None:
    grouped: Dict[str, List[Dict[str, Any]]] = {}
    for row in comparisons:
        grouped.setdefault(row["query"], []).append(row)

    ordered_queries = sorted(
        grouped,
        key=lambda query: min(item["delta_primary"] for item in grouped[query]),
    )

    print(f"Detailed Examples (top {top_queries} queries)")
    print("=" * 80)
    for query in ordered_queries[:top_queries]:
        rows = sorted(grouped[query], key=lambda item: item["backup_rank_eval"])
        print(f"\n## {query}")
        print(f"affected_docs={len(rows)} | delta_primary={rows[0]['delta_primary']:+.6f}")
        for row in rows[:docs_per_query]:
            current_row = row["current_row"]
            backup_row = row["backup_row"]
            print(
                f"  - spu={row['spu_id']} "
                f"eval_current={row['current_rank_eval']} eval_backup={row['backup_rank_eval']} "
                f"coarse_current={current_row.get('_coarse_rank') if current_row else None} "
                f"coarse_backup={backup_row.get('_coarse_rank') if backup_row else None}"
            )
            if current_row and backup_row:
                print(
                    "    image_knn "
                    f"{backup_row.get('image_knn_score')} -> {current_row.get('image_knn_score')} | "
                    "text_knn "
                    f"{backup_row.get('text_knn_score')} -> {current_row.get('text_knn_score')} | "
                    "text_score "
                    f"{backup_row.get('text_score')} -> {current_row.get('text_score')} | "
                    "es_score "
                    f"{backup_row.get('es_score')} -> {current_row.get('es_score')} | "
                    "coarse_score "
                    f"{backup_row.get('coarse_score')} -> {current_row.get('coarse_score')}"
                )
            else:
                print(
                    f"    present_current={current_row is not None} "
                    f"present_backup={backup_row is not None}"
                )


def main() -> None:
    parser = argparse.ArgumentParser(description="Analyze coarse-score component regressions")
    parser.add_argument("--current-report", required=True)
    parser.add_argument("--backup-report", required=True)
    parser.add_argument("--current-index", required=True)
    parser.add_argument("--backup-index", required=True)
    parser.add_argument("--tenant-id", default="163")
    parser.add_argument("--rank-gap-threshold", type=int, default=5)
    parser.add_argument("--scan-depth", type=int, default=20)
    parser.add_argument("--detail-queries", type=int, default=6)
    parser.add_argument("--detail-docs-per-query", type=int, default=3)
    args = parser.parse_args()

    logging.basicConfig(level=logging.WARNING)

    current_report = _load_report(args.current_report)
    backup_report = _load_report(args.backup_report)
    regressed = _collect_regressed_docs(
        current_report=current_report,
        backup_report=backup_report,
        rank_gap_threshold=args.rank_gap_threshold,
        scan_depth=args.scan_depth,
    )

    searcher = _build_searcher()
    comparisons: List[Dict[str, Any]] = []

    for query, rows in regressed.items():
        current_by_doc, _ = _run_query(
            searcher,
            query=query,
            tenant_id=args.tenant_id,
            index_name=args.current_index,
        )
        backup_by_doc, _ = _run_query(
            searcher,
            query=query,
            tenant_id=args.tenant_id,
            index_name=args.backup_index,
        )
        for row in rows:
            comparisons.append(
                {
                    **row,
                    "current_row": current_by_doc.get(row["spu_id"]),
                    "backup_row": backup_by_doc.get(row["spu_id"]),
                }
            )

    _summarize_rows(comparisons)
    _print_query_examples(
        comparisons,
        top_queries=args.detail_queries,
        docs_per_query=args.detail_docs_per_query,
    )


if __name__ == "__main__":
    main()