from __future__ import annotations

from typing import Any

import pandas as pd

from featureleak.report import LeakIssue
from featureleak.utils.schema import infer_schema


def _safe_overlap_stats(train_values: pd.Series, test_values: pd.Series) -> dict[str, Any]:
    try:
        train_set = set(train_values.dropna().astype(str).tolist())
        test_set = set(test_values.dropna().astype(str).tolist())

        overlap = train_set & test_set
        union = train_set | test_set

        return {
            "n_train_entities": len(train_set),
            "n_test_entities": len(test_set),
            "n_overlap_entities": len(overlap),
            "train_overlap_ratio": len(overlap) / max(len(train_set), 1),
            "test_overlap_ratio": len(overlap) / max(len(test_set), 1),
            "jaccard_overlap": len(overlap) / max(len(union), 1),
            "sample_overlap_entities": sorted(list(overlap))[:10],
        }
    except Exception:
        return {
            "n_train_entities": 0,
            "n_test_entities": 0,
            "n_overlap_entities": 0,
            "train_overlap_ratio": 0.0,
            "test_overlap_ratio": 0.0,
            "jaccard_overlap": 0.0,
            "sample_overlap_entities": [],
        }


def _candidate_entity_columns(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
) -> list[str]:
    """
    Fall back to inferred identifier-like columns if user does not provide entity columns.
    """
    shared = [c for c in train_df.columns if c in test_df.columns]
    if not shared:
        return []

    combined = pd.concat(
        [train_df[shared].head(5000), test_df[shared].head(5000)],
        axis=0,
        ignore_index=True,
    )
    schema = infer_schema(combined)
    return list(schema.get("identifier_columns", []))


def detect_entity_leakage(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    entity_cols: list[str] | None = None,
    target: str | None = None,
    overlap_threshold: float = 0.50,
    severe_overlap_threshold: float = 0.90,
) -> list[LeakIssue]:
    """
    Detect overlap of entity/group identifiers across train and test.

    Examples:
    - same customer in both train and test
    - same patient in both splits
    - same store/account/device across splits

    This is only a leakage risk if the user's intended evaluation requires
    generalization to unseen entities, but it is worth flagging clearly.
    """
    issues: list[LeakIssue] = []

    if train_df.empty or test_df.empty:
        return issues

    if entity_cols is None or len(entity_cols) == 0:
        entity_cols = _candidate_entity_columns(train_df, test_df)

    if not entity_cols:
        return issues

    valid_entity_cols = [c for c in entity_cols if c in train_df.columns and c in test_df.columns]
    if not valid_entity_cols:
        return issues

    for col in valid_entity_cols:
        stats = _safe_overlap_stats(train_df[col], test_df[col])

        train_overlap_ratio = float(stats["train_overlap_ratio"])
        test_overlap_ratio = float(stats["test_overlap_ratio"])
        jaccard_overlap = float(stats["jaccard_overlap"])
        n_overlap_entities = int(stats["n_overlap_entities"])

        if n_overlap_entities == 0:
            continue

        max_overlap = max(train_overlap_ratio, test_overlap_ratio)

        if max_overlap >= severe_overlap_threshold:
            severity = "high"
        elif max_overlap >= overlap_threshold:
            severity = "medium"
        else:
            continue

        issues.append(
            LeakIssue(
                severity=severity,
                check="entity_leakage",
                title="Entities overlap across train and test",
                summary=(
                    f"Entity column `{col}` has overlapping values across train and test. "
                    f"This can inflate evaluation if you need unseen-entity generalization."
                ),
                affected_columns=[col],
                details={
                    "n_train_entities": stats["n_train_entities"],
                    "n_test_entities": stats["n_test_entities"],
                    "n_overlap_entities": n_overlap_entities,
                    "train_overlap_ratio": round(train_overlap_ratio, 6),
                    "test_overlap_ratio": round(test_overlap_ratio, 6),
                    "jaccard_overlap": round(jaccard_overlap, 6),
                    "sample_overlap_entities": stats["sample_overlap_entities"],
                },
                suggested_fix=(
                    "If your deployment setting requires generalization to unseen entities, "
                    "use a group-based split such as GroupKFold or a holdout split by entity."
                ),
            )
        )

    return issues