from __future__ import annotations

from typing import Any

import pandas as pd

from featureleak.report import LeakIssue
from featureleak.utils.hashing import stable_row_hashes


def _safe_overlap_keys(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    key_cols: list[str],
) -> dict[str, Any]:
    if not key_cols:
        return {
            "n_train_keys": 0,
            "n_test_keys": 0,
            "n_overlap_keys": 0,
            "train_overlap_ratio": 0.0,
            "test_overlap_ratio": 0.0,
            "sample_overlap_keys": [],
        }

    train_keys = (
        train_df[key_cols]
        .astype("object")
        .where(train_df[key_cols].notna(), "__NaN__")
        .astype(str)
        .agg("||".join, axis=1)
    )
    test_keys = (
        test_df[key_cols]
        .astype("object")
        .where(test_df[key_cols].notna(), "__NaN__")
        .astype(str)
        .agg("||".join, axis=1)
    )

    train_unique = pd.Index(train_keys.drop_duplicates())
    test_unique = pd.Index(test_keys.drop_duplicates())
    overlap = train_unique.intersection(test_unique)

    return {
        "n_train_keys": int(len(train_unique)),
        "n_test_keys": int(len(test_unique)),
        "n_overlap_keys": int(len(overlap)),
        "train_overlap_ratio": float(len(overlap) / max(len(train_unique), 1)),
        "test_overlap_ratio": float(len(overlap) / max(len(test_unique), 1)),
        "sample_overlap_keys": sorted(map(str, overlap[:10].tolist())),
    }


def detect_contamination(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    target: str | None = None,
    overlap_threshold: float = 0.0,
    key_columns: list[str] | None = None,
) -> list[LeakIssue]:
    """
    Detect train/test contamination via overlapping row hashes.

    Behavior:
    - exact overlap check excluding target
    - optional business-key overlap for user-supplied key columns
    """
    issues: list[LeakIssue] = []

    if train_df.empty or test_df.empty:
        return issues

    exclude_cols = [target] if target is not None and target in train_df.columns else []

    train_hashes = stable_row_hashes(train_df, exclude_cols=exclude_cols)
    test_hashes = stable_row_hashes(test_df, exclude_cols=exclude_cols)

    train_unique_hashes = pd.Index(train_hashes.drop_duplicates())
    test_unique_hashes = pd.Index(test_hashes.drop_duplicates())
    overlap_hashes = train_unique_hashes.intersection(test_unique_hashes)

    if len(overlap_hashes) > 0:
        train_overlap_mask = train_hashes.isin(overlap_hashes)
        test_overlap_mask = test_hashes.isin(overlap_hashes)

        n_train_overlap_rows = int(train_overlap_mask.sum())
        n_test_overlap_rows = int(test_overlap_mask.sum())

        overlap_rate_train = n_train_overlap_rows / max(len(train_df), 1)
        overlap_rate_test = n_test_overlap_rows / max(len(test_df), 1)
        overlap_rate = max(overlap_rate_train, overlap_rate_test)

        if overlap_rate >= overlap_threshold:
            severity = "high"
            if overlap_rate >= 0.10:
                severity = "critical"

            issues.append(
                LeakIssue(
                    severity=severity,
                    check="train_test_contamination",
                    title="Overlapping rows across train and test",
                    summary=(
                        f"{len(overlap_hashes)} overlapping row signatures were found across "
                        f"train and test after excluding the target column."
                    ),
                    affected_columns=[],
                    details={
                        "n_overlap_signatures": int(len(overlap_hashes)),
                        "n_train_overlap_rows": n_train_overlap_rows,
                        "n_test_overlap_rows": n_test_overlap_rows,
                        "train_overlap_rate": round(overlap_rate_train, 6),
                        "test_overlap_rate": round(overlap_rate_test, 6),
                        "excluded_columns": exclude_cols,
                    },
                    suggested_fix=(
                        "Deduplicate before splitting, or rebuild the split logic to ensure "
                        "the same examples do not appear in both train and test."
                    ),
                )
            )

    valid_key_columns = [
        c for c in (key_columns or []) if c in train_df.columns and c in test_df.columns
    ]

    if valid_key_columns:
        key_stats = _safe_overlap_keys(train_df, test_df, valid_key_columns)
        if key_stats["n_overlap_keys"] > 0:
            max_overlap = max(
                key_stats["train_overlap_ratio"],
                key_stats["test_overlap_ratio"],
            )

            if max_overlap >= max(overlap_threshold, 0.05):
                severity = "medium"
                if max_overlap >= 0.30:
                    severity = "high"

                issues.append(
                    LeakIssue(
                        severity=severity,
                        check="train_test_contamination",
                        title="Business-key overlap across train and test",
                        summary=(
                            f"User-specified key columns {valid_key_columns} overlap across "
                            f"train and test, which may indicate split contamination."
                        ),
                        affected_columns=valid_key_columns,
                        details=key_stats,
                        suggested_fix=(
                            "Check whether these keys should be unique to one split. If so, "
                            "split on the business key rather than on rows."
                        ),
                    )
                )

    return issues