"""Evaluation functions for measuring identity resolution quality."""
from __future__ import annotations

from dataclasses import dataclass
from itertools import combinations
from typing import Any

import json

from .reconcile import ReconcileResult, reconcile
from .source import Source
from .spec import Spec


@dataclass
class EvaluateResult:
    """Precision, recall, and F1 metrics for an identity resolution run."""

    precision: float
    recall: float
    f1: float
    true_positives: int
    false_positives: int
    false_negatives: int
    predicted_pairs: int
    ground_truth_pairs: int
    predicted_clusters: int
    ground_truth_clusters: int
    total_records: int

    def summary(self) -> str:
        """Return a human-readable summary of the evaluation."""
        lines = [
            "Evaluation Results",
            "=" * 40,
            f"  Precision:          {self.precision:.4f}",
            f"  Recall:             {self.recall:.4f}",
            f"  F1 Score:           {self.f1:.4f}",
            "",
            f"  True Positives:     {self.true_positives}",
            f"  False Positives:    {self.false_positives}",
            f"  False Negatives:    {self.false_negatives}",
            "",
            f"  Predicted Pairs:    {self.predicted_pairs}",
            f"  Ground Truth Pairs: {self.ground_truth_pairs}",
            f"  Predicted Clusters: {self.predicted_clusters}",
            f"  GT Clusters:        {self.ground_truth_clusters}",
            f"  Total Records:      {self.total_records}",
        ]
        return "\n".join(lines)

    def __repr__(self) -> str:
        return (
            f"EvaluateResult(precision={self.precision:.4f}, "
            f"recall={self.recall:.4f}, f1={self.f1:.4f})"
        )


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _pairs_from_clusters(
    clusters: list[set[tuple[str, str]]],
) -> set[tuple[tuple[str, str], tuple[str, str]]]:
    """Convert list of member sets to a canonical pairwise set.

    Each pair is sorted so (A, B) and (B, A) produce the same tuple.
    """
    pairs: set[tuple[tuple[str, str], tuple[str, str]]] = set()
    for cluster in clusters:
        for a, b in combinations(sorted(cluster), 2):
            pairs.add((a, b))
    return pairs


def _clusters_from_result(
    result: ReconcileResult,
    entity_map: dict[str, tuple[str, str]],
) -> list[set[tuple[str, str]]]:
    """Convert UUID-based clusters to (source_name, external_id) clusters."""
    clusters: list[set[tuple[str, str]]] = []
    for uuid_cluster in result.clusters:
        member_set: set[tuple[str, str]] = set()
        for uid in uuid_cluster:
            if uid in entity_map:
                member_set.add(entity_map[uid])
        if len(member_set) >= 2:
            clusters.append(member_set)
    return clusters


def _parse_ground_truth(
    ground_truth: Any,
) -> list[set[tuple[str, str]]]:
    """Normalize ground truth from dict or DataFrame to list of member sets.

    Accepts:
        - dict: {true_entity_id: [(source_name, external_id), ...]}
        - pandas DataFrame with columns [record_id, source_name, true_entity_id]
    """
    # DataFrame
    try:
        import pandas as pd

        if isinstance(ground_truth, pd.DataFrame):
            required = {"record_id", "source_name", "true_entity_id"}
            missing = required - set(ground_truth.columns)
            if missing:
                raise ValueError(
                    f"Ground truth DataFrame missing columns: {sorted(missing)}"
                )
            clusters: dict[str, set[tuple[str, str]]] = {}
            for _, row in ground_truth.iterrows():
                eid = str(row["true_entity_id"])
                clusters.setdefault(eid, set()).add(
                    (str(row["source_name"]), str(row["record_id"]))
                )
            return [c for c in clusters.values() if len(c) >= 2]
    except ImportError:
        pass

    # Dict format
    if isinstance(ground_truth, dict):
        result: list[set[tuple[str, str]]] = []
        for _eid, members in ground_truth.items():
            member_set = {(str(s), str(r)) for s, r in members}
            if len(member_set) >= 2:
                result.append(member_set)
        return result

    raise TypeError(
        f"ground_truth must be a dict or pandas DataFrame, got {type(ground_truth).__name__}"
    )


def _build_holdout_ground_truth(
    sources: list[Source],
    field: str,
    min_size: int = 2,
) -> list[set[tuple[str, str]]]:
    """Build ground truth clusters by grouping records that share a holdout field value."""
    groups: dict[str, set[tuple[str, str]]] = {}

    for source in sources:
        if source.primary_key is None:
            raise ValueError(
                f"Source '{source.name}' has no primary_key set. "
                "Holdout evaluation requires primary_key on all sources."
            )
        for row in source.iter_rows():
            val = row.get(field, "")
            if not val or val.lower() in ("", "none", "null", "nan"):
                continue
            groups.setdefault(val, set()).add((source.name, row[source.primary_key]))

    clusters = [members for members in groups.values() if len(members) >= min_size]
    if not clusters:
        raise ValueError(
            f"No valid ground truth clusters found for holdout field '{field}'. "
            f"Need at least {min_size} records sharing the same field value."
        )
    return clusters


def _remove_field_from_spec(spec: Spec, field: str) -> Spec:
    """Return a new Spec with the holdout field removed from rules, FS fields, and blocking."""
    # Deep-copy via JSON round-trip (parsed is already a plain dict from Rust)
    parsed = json.loads(json.dumps(spec.parsed))

    # Remove from rules (rules-based path)
    if "rules" in parsed:
        parsed["rules"] = [
            r for r in parsed["rules"]
            if r.get("field") != field
        ]
        if not parsed["rules"]:
            del parsed["rules"]

    # Remove from decision.scoring.fields (Fellegi-Sunter path)
    scoring = parsed.get("decision", {}).get("scoring", {})
    if "fields" in scoring:
        scoring["fields"] = [
            f for f in scoring["fields"]
            if f.get("name") != field
        ]

    # Remove from blocking keys
    blocking = parsed.get("blocking", {})
    if "keys" in blocking:
        new_keys = []
        for key_list in blocking["keys"]:
            filtered = [k for k in key_list if k != field]
            if filtered:
                new_keys.append(filtered)
        if new_keys:
            blocking["keys"] = new_keys
        else:
            del parsed["blocking"]

    # JSON is valid YAML - the Rust parser handles it fine
    return Spec.from_string(json.dumps(parsed))


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def evaluate(
    result: ReconcileResult,
    sources: list[Source],
    ground_truth: Any,
) -> EvaluateResult:
    """Evaluate identity resolution quality against labeled ground truth.

    Args:
        result: A ``ReconcileResult`` with ``_entity_map`` attached
            (from ``reconcile(..., return_entity_map=True)``).
        sources: The original sources used for reconciliation.
        ground_truth: Labeled data in one of two formats:
            - dict: ``{entity_id: [(source_name, external_id), ...]}``
            - pandas DataFrame with columns ``[record_id, source_name, true_entity_id]``

    Returns:
        EvaluateResult with precision, recall, F1, and supporting counts.
    """
    entity_map = getattr(result, "_entity_map", None)
    if entity_map is None:
        raise ValueError(
            "ReconcileResult has no _entity_map. "
            "Call reconcile(..., return_entity_map=True) first."
        )

    # Build predicted and ground truth pair sets
    predicted_clusters = _clusters_from_result(result, entity_map)
    gt_clusters = _parse_ground_truth(ground_truth)

    # Collect all records present in ground truth
    gt_records: set[tuple[str, str]] = set()
    for cluster in gt_clusters:
        gt_records.update(cluster)

    # Filter predicted clusters to only include records in ground truth
    filtered_predicted: list[set[tuple[str, str]]] = []
    for cluster in predicted_clusters:
        filtered = cluster & gt_records
        if len(filtered) >= 2:
            filtered_predicted.append(filtered)

    predicted_pairs = _pairs_from_clusters(filtered_predicted)
    gt_pairs = _pairs_from_clusters(gt_clusters)

    tp = len(predicted_pairs & gt_pairs)
    fp = len(predicted_pairs - gt_pairs)
    fn = len(gt_pairs - predicted_pairs)

    precision = tp / (tp + fp) if (tp + fp) > 0 else 1.0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 1.0
    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) > 0 else 0.0

    total_records = sum(len(c) for c in result.clusters)

    return EvaluateResult(
        precision=precision,
        recall=recall,
        f1=f1,
        true_positives=tp,
        false_positives=fp,
        false_negatives=fn,
        predicted_pairs=len(predicted_pairs),
        ground_truth_pairs=len(gt_pairs),
        predicted_clusters=len(filtered_predicted),
        ground_truth_clusters=len(gt_clusters),
        total_records=total_records,
    )


def evaluate_holdout(
    sources: list[Source],
    spec: Spec,
    holdout_field: str,
    *,
    min_cluster_size: int = 2,
) -> EvaluateResult:
    """Evaluate resolution quality using holdout methodology.

    Holds out a strong identifier (e.g. email), re-runs resolution without it,
    and scores the result against ground truth clusters built from the held-out
    field.

    Args:
        sources: Data sources to reconcile.
        spec: The identity spec.
        holdout_field: Field to hold out (e.g. "email"). Must be present in
            the spec's rules, FS fields, or blocking keys.
        min_cluster_size: Minimum records sharing a holdout value to form a
            ground truth cluster. Default 2.

    Returns:
        EvaluateResult with precision, recall, F1 measured against the
        constructed ground truth.
    """
    # Step 1: Build ground truth from the holdout field
    gt_clusters = _build_holdout_ground_truth(
        sources, holdout_field, min_size=min_cluster_size
    )

    # Step 2: Remove the holdout field from the spec
    modified_spec = _remove_field_from_spec(spec, holdout_field)

    # Step 3: Re-run reconciliation without the holdout field
    result = reconcile(sources, modified_spec, return_entity_map=True)

    # Step 4: Convert ground truth to dict format for evaluate()
    gt_dict: dict[int, list[tuple[str, str]]] = {}
    for i, cluster in enumerate(gt_clusters):
        gt_dict[i] = list(cluster)

    return evaluate(result, sources, gt_dict)
