"""NuExtract benchmark comparing against GLiNER2 + deterministic parsing.

Evaluates NuExtract 1.5 Smol on hierarchical extraction for long filings
and transcripts, measuring incremental correctness, CPU latency, and memory.

Reports per-document-type incremental value to determine which document
classes benefit from NuExtract supplementation.

Requirement: 6.6
"""
from __future__ import annotations

import logging
from collections import defaultdict
from typing import Any

from services.intelligence_pipeline_v3.nuextract.adapter import NuExtractAdapter
from services.intelligence_pipeline_v3.nuextract.models import (
    BenchmarkReport,
    IncrementalValueReport,
    NuExtractResult,
    PromotionGate,
)
from services.intelligence_pipeline_v3.nuextract.promotion import PromotionEvaluator

logger = logging.getLogger(__name__)


class GoldDocument:
    """A document from the gold corpus with ground-truth labels."""

    def __init__(
        self,
        text: str,
        document_type: str,
        gold_fields: dict[str, Any],
        schema: dict[str, Any],
        document_id: str = "",
    ) -> None:
        self.text = text
        self.document_type = document_type
        self.gold_fields = gold_fields
        self.schema = schema
        self.document_id = document_id


class GLiNERResult:
    """Simulated result from GLiNER2 + deterministic parsing."""

    def __init__(
        self,
        fields: dict[str, Any],
        latency_ms: float = 0.0,
        memory_mb: float = 0.0,
    ) -> None:
        self.fields = fields
        self.latency_ms = latency_ms
        self.memory_mb = memory_mb


class NuExtractBenchmark:
    """Benchmark comparing NuExtract vs GLiNER2 + deterministic parsing.

    Evaluates per-document-type to determine where NuExtract adds value.

    Parameters
    ----------
    adapter
        NuExtractAdapter instance (test_mode or production).
    gate
        Promotion gate thresholds for deciding promotion.
    """

    def __init__(
        self,
        adapter: NuExtractAdapter | None = None,
        gate: PromotionGate | None = None,
    ) -> None:
        self._adapter = adapter or NuExtractAdapter(test_mode=True)
        self._gate = gate or PromotionGate()
        self._evaluator = PromotionEvaluator(self._gate)

    async def evaluate_against_gliner(
        self,
        documents: list[GoldDocument],
        gliner_results: list[GLiNERResult],
    ) -> BenchmarkReport:
        """Run full benchmark comparing NuExtract vs GLiNER2 + deterministic parsing.

        Parameters
        ----------
        documents
            Gold corpus documents with ground-truth labels.
        gliner_results
            Pre-computed GLiNER2 + deterministic parsing results for each document.

        Returns
        -------
        BenchmarkReport
            Full benchmark report with per-type results and promotion decisions.
        """
        if len(documents) != len(gliner_results):
            raise ValueError(
                f"Document count ({len(documents)}) must match "
                f"GLiNER result count ({len(gliner_results)})"
            )

        # Group by document type
        by_type: dict[str, list[tuple[GoldDocument, GLiNERResult]]] = defaultdict(list)
        for doc, gliner in zip(documents, gliner_results):
            by_type[doc.document_type].append((doc, gliner))

        # Evaluate each document type
        reports: list[IncrementalValueReport] = []
        for doc_type, pairs in by_type.items():
            report = await self._evaluate_type(doc_type, pairs)
            reports.append(report)

        # Determine promotions
        promoted_types: list[str] = []
        for report in reports:
            if self._evaluator.evaluate(report):
                report.promoted = True
                promoted_types.append(report.document_type)

        # Compute overall metrics
        total_docs = len(documents)
        overall_nuextract_f1 = 0.0
        overall_gliner_f1 = 0.0
        if reports:
            weighted_nu = sum(r.nuextract_f1 * r.sample_count for r in reports)
            weighted_gl = sum(r.gliner_f1 * r.sample_count for r in reports)
            overall_nuextract_f1 = weighted_nu / total_docs if total_docs > 0 else 0.0
            overall_gliner_f1 = weighted_gl / total_docs if total_docs > 0 else 0.0

        return BenchmarkReport(
            reports=reports,
            gate=self._gate,
            promoted_types=promoted_types,
            overall_nuextract_f1=overall_nuextract_f1,
            overall_gliner_f1=overall_gliner_f1,
            overall_delta=overall_nuextract_f1 - overall_gliner_f1,
            total_documents=total_docs,
        )

    async def _evaluate_type(
        self,
        doc_type: str,
        pairs: list[tuple[GoldDocument, GLiNERResult]],
    ) -> IncrementalValueReport:
        """Evaluate NuExtract vs GLiNER for a single document type."""
        nuextract_scores: list[float] = []
        gliner_scores: list[float] = []
        nuextract_latencies: list[float] = []
        nuextract_memories: list[float] = []
        gliner_latencies: list[float] = []
        gliner_memories: list[float] = []

        for doc, gliner_result in pairs:
            # Run NuExtract extraction
            nu_result = await self._adapter.extract(
                text=doc.text,
                schema=doc.schema,
                document_type=doc.document_type,
            )

            # Compute F1 for NuExtract
            nu_f1 = self._compute_field_f1(nu_result, doc.gold_fields)
            nuextract_scores.append(nu_f1)
            nuextract_latencies.append(nu_result.latency_ms)
            nuextract_memories.append(nu_result.memory_mb)

            # Compute F1 for GLiNER
            gl_f1 = self._compute_extraction_f1(gliner_result.fields, doc.gold_fields)
            gliner_scores.append(gl_f1)
            gliner_latencies.append(gliner_result.latency_ms)
            gliner_memories.append(gliner_result.memory_mb)

        # Aggregate metrics
        n = len(pairs)
        avg_nu_f1 = sum(nuextract_scores) / n if n > 0 else 0.0
        avg_gl_f1 = sum(gliner_scores) / n if n > 0 else 0.0
        p95_nu_latency = _percentile(nuextract_latencies, 95)
        p95_gl_latency = _percentile(gliner_latencies, 95)
        max_nu_memory = max(nuextract_memories) if nuextract_memories else 0.0
        max_gl_memory = max(gliner_memories) if gliner_memories else 0.0

        return IncrementalValueReport(
            document_type=doc_type,
            gliner_f1=avg_gl_f1,
            nuextract_f1=avg_nu_f1,
            delta=avg_nu_f1 - avg_gl_f1,
            nuextract_latency_ms=p95_nu_latency,
            nuextract_memory_mb=max_nu_memory,
            gliner_latency_ms=p95_gl_latency,
            gliner_memory_mb=max_gl_memory,
            sample_count=n,
            promoted=False,
        )

    def _compute_field_f1(
        self, result: NuExtractResult, gold: dict[str, Any]
    ) -> float:
        """Compute F1 score for NuExtract result against gold labels."""
        if not gold:
            return 1.0 if not result.fields else 0.0

        extracted_fields = {
            f.name: f.value for f in result.fields if f.value is not None
        }
        return self._compute_extraction_f1(extracted_fields, gold)

    def _compute_extraction_f1(
        self, predicted: dict[str, Any], gold: dict[str, Any]
    ) -> float:
        """Compute field-level F1 between predicted and gold extractions."""
        if not gold and not predicted:
            return 1.0
        if not gold or not predicted:
            return 0.0

        gold_set = set(gold.keys())
        pred_set = set(predicted.keys())

        # True positives: predicted fields that match gold (key present AND value matches)
        tp = 0
        for key in gold_set & pred_set:
            if self._values_match(predicted[key], gold[key]):
                tp += 1

        precision = tp / len(pred_set) if pred_set else 0.0
        recall = tp / len(gold_set) if gold_set else 0.0

        if precision + recall == 0:
            return 0.0
        return 2 * precision * recall / (precision + recall)

    def _values_match(self, predicted: Any, gold: Any) -> bool:
        """Check if a predicted value matches gold (with tolerance)."""
        if predicted is None:
            return gold is None
        if gold is None:
            return False

        # String comparison (case-insensitive, trimmed)
        if isinstance(gold, str) and isinstance(predicted, str):
            return predicted.strip().lower() == gold.strip().lower()

        # Numeric comparison with tolerance
        if isinstance(gold, (int, float)) and isinstance(predicted, (int, float)):
            if gold == 0:
                return abs(predicted) < 1e-6
            return abs(predicted - gold) / abs(gold) < 0.05

        # Dict comparison (recursive for hierarchical)
        if isinstance(gold, dict) and isinstance(predicted, dict):
            if not gold:
                return not predicted
            matches = sum(
                1
                for k in gold
                if k in predicted and self._values_match(predicted[k], gold[k])
            )
            return matches / len(gold) >= 0.5

        # Fallback: equality
        return predicted == gold


def _percentile(values: list[float], pct: int) -> float:
    """Compute a percentile from a list of values."""
    if not values:
        return 0.0
    sorted_vals = sorted(values)
    idx = int(len(sorted_vals) * pct / 100)
    idx = min(idx, len(sorted_vals) - 1)
    return sorted_vals[idx]
