"""NCM retrieval evaluation V2 (sections 1, 13).

Same sampling as V1 (seed 42, per-segment) with each sample now quality-classified:
  GOLD   - GTIN product or cross-tenant agreement on the same fingerprint/NCM
  SILVER - single-tenant valid NCM, no independent confirmation
  NOISY  - commercial description with < 2 informative tokens (not pursued)
Retrieval measured separately: OFFICIAL_ONLY (hierarchical enriched corpus),
KB_ONLY (neighbor bridge), HYBRID; recall@1/5/10/20/50 + MRR, per class and segment.
"""
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from sqlalchemy import select

from central.app.ai.retrieval import (EnrichedNcmCorpus, KbNeighborIndex, commercial_tokens, derive_abbreviation_lexicon,
                                      hierarchical_candidates, hybrid_candidates, kb_neighbor_candidates)
from central.app.db import SessionLocal
from central.app.models import Dataset, KbObservation, KbProduct, NcmEntry

SEGMENTS = {"CONSTRUCAO": "construcao", "MMMODAS": "moda", "RINCAO": "restaurante", "UNIAO": "mercado", "TESTESOFICIAIS": "base_testes"}
KS = (1, 5, 10, 20, 50)


def _samples(session, per_segment=40, seed=42):
    import random
    dataset = session.scalar(select(Dataset).where(Dataset.kind == "NCM", Dataset.status == "ACTIVE"))
    valid_ncm = set(session.scalars(select(NcmEntry.code).where(NcmEntry.dataset_id == dataset.id)).all())
    rows = session.execute(
        select(KbProduct, KbObservation).join(KbObservation, KbObservation.kb_product_id == KbProduct.id)
    ).all()
    per_product_tenants = {}
    for product, observation in rows:
        per_product_tenants.setdefault(product.id, set()).add(observation.tenant_id)
    by_segment = {}
    for product, observation in rows:
        ncm = (observation.observed_fields or {}).get("ncm")
        if ncm and ncm in valid_ncm and product.normalized_description:
            by_segment.setdefault(SEGMENTS.get(observation.database_label, observation.database_label), set()).add(
                (product.id, product.normalized_description, ncm, bool(product.gtin), len(per_product_tenants[product.id]) >= 2)
            )
    rng = random.Random(seed)
    samples = []
    for segment, values in sorted(by_segment.items()):
        picked = sorted(values) if len(values) <= per_segment else rng.sample(sorted(values), per_segment)
        for product_id, description, ncm, has_gtin, cross_tenant in picked:
            quality = "NOISY" if len(commercial_tokens(description)) < 2 else ("GOLD" if (has_gtin or cross_tenant) else "SILVER")
            samples.append({"segment": segment, "product_id": product_id, "description": description, "true_ncm": ncm, "quality": quality})
    return samples


def _recall_mrr(runs, samples):
    """runs: {mode: [ranked_codes_per_sample]}"""
    out = {}
    for mode, ranked_lists in runs.items():
        metrics = {f"recall_at_{k}": 0 for k in KS} | {"mrr": 0.0}
        total = 0
        for sample, ranked in zip(samples, ranked_lists):
            codes = [row[0] for row in ranked]
            total += 1
            for k in KS:
                if sample["true_ncm"] in codes[:k]:
                    metrics[f"recall_at_{k}"] += 1
            if sample["true_ncm"] in codes:
                metrics["mrr"] += 1.0 / (codes.index(sample["true_ncm"]) + 1)
        if total:
            for key in metrics:
                metrics[key] = round(metrics[key] / total, 4)
        out[mode] = {"total": total, **metrics}
    return out


def main(per_segment=40, output="docs/NCM_RETRIEVAL_EVALUATION_V2.json") -> dict:
    session = SessionLocal()
    samples = _samples(session, per_segment)
    corpus = EnrichedNcmCorpus.build(session)
    index = KbNeighborIndex.build(session)
    lexicon = derive_abbreviation_lexicon(index.corpus(), min_support=3)
    runs = {"OFFICIAL_ONLY": [], "KB_ONLY": [], "HYBRID": []}
    for sample in samples:
        # leave-one-out: the queried product never retrieves itself from the KB
        runs["OFFICIAL_ONLY"].append(hierarchical_candidates(corpus, sample["description"], k=50))
        runs["KB_ONLY"].append(kb_neighbor_candidates(index, sample["description"], k=50, exclude_product_id=sample["product_id"]))
        runs["HYBRID"].append(hybrid_candidates(corpus, index, sample["description"], k=50, lexicon=lexicon, exclude_product_id=sample["product_id"]))
    overall = _recall_mrr(runs, samples)
    by_quality = {quality: _recall_mrr({mode: [row for row, sample in zip(list_rows, samples) if sample["quality"] == quality] for mode, list_rows in runs.items()},
                                       [sample for sample in samples if sample["quality"] == quality]) for quality in ("GOLD", "SILVER", "NOISY")}
    by_segment = {segment: _recall_mrr({mode: [row for row, sample in zip(list_rows, samples) if sample["segment"] == segment] for mode, list_rows in runs.items()},
                                        [sample for sample in samples if sample["segment"] == segment]) for segment in sorted({sample["segment"] for sample in samples})}
    report = {
        "artifact": "NCM_RETRIEVAL_EVALUATION_V2",
        "eval_total": len(samples),
        "quality_counts": {quality: sum(1 for sample in samples if sample["quality"] == quality) for quality in ("GOLD", "SILVER", "NOISY")},
        "lexicon_size": len(lexicon),
        "lexicon_sample": dict(sorted(lexicon.items())[:40]),
        "overall": overall,
        "by_quality": by_quality,
        "by_segment": by_segment,
        "llm_liberation_gate": {"gold_recall_at_20_min": 0.80, "gold_recall_at_50_preferred": 0.90},
        "cost_usd": 0.0,
    }
    Path(output).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
    session.close()
    return report


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--per-segment", type=int, default=40)
    parser.add_argument("--output", default="docs/NCM_RETRIEVAL_EVALUATION_V2.json")
    args = parser.parse_args()
    print(json.dumps(main(args.per_segment, args.output), indent=2))