from __future__ import annotations

from typing import Any

import pandas as pd
from pandas.api.types import is_numeric_dtype

from featureleak.report import LeakIssue


def _null_target_association_numeric_target(
    is_null: pd.Series,
    target: pd.Series,
) -> dict[str, Any]:
    """
    Association metrics when target is numeric.
    """
    out: dict[str, Any] = {
        "metric_name": None,
        "metric_value": None,
        "null_target_mean": None,
        "non_null_target_mean": None,
    }

    try:
        mask = is_null.astype(bool)
        numeric_target = pd.to_numeric(target, errors="coerce")

        if numeric_target.notna().sum() == 0:
            return out

        null_mean = numeric_target[mask].mean() if mask.any() else None
        non_null_mean = numeric_target[~mask].mean() if (~mask).any() else None

        corr = pd.Series(mask.astype(int), index=target.index).corr(numeric_target)

        out["metric_name"] = "abs_null_target_corr"
        out["metric_value"] = None if pd.isna(corr) else float(abs(corr))
        out["null_target_mean"] = None if pd.isna(null_mean) else float(null_mean)
        out["non_null_target_mean"] = None if pd.isna(non_null_mean) else float(non_null_mean)
        return out
    except Exception:
        return out


def _null_target_association_categorical_target(
    is_null: pd.Series,
    target: pd.Series,
) -> dict[str, Any]:
    """
    Association metrics when target is non-numeric / categorical-like.
    """
    out: dict[str, Any] = {
        "metric_name": None,
        "metric_value": None,
        "null_positive_rate": None,
        "non_null_positive_rate": None,
    }

    try:
        mask = is_null.astype(bool)
        safe_target = target.astype("object").where(target.notna(), "__NaN__").astype(str)

        dist_null = safe_target[mask].value_counts(normalize=True, dropna=False)
        dist_non_null = safe_target[~mask].value_counts(normalize=True, dropna=False)

        all_labels = set(dist_null.index) | set(dist_non_null.index)
        total_variation = 0.5 * sum(abs(dist_null.get(lbl, 0.0) - dist_non_null.get(lbl, 0.0)) for lbl in all_labels)

        out["metric_name"] = "target_distribution_shift"
        out["metric_value"] = float(total_variation)
        return out
    except Exception:
        return out


def _null_rate_shift(train_series: pd.Series, test_series: pd.Series) -> float | None:
    try:
        return float(abs(train_series.isna().mean() - test_series.isna().mean()))
    except Exception:
        return None


def detect_missingness_leakage(
    df: pd.DataFrame,
    target: str,
    train_df: pd.DataFrame | None = None,
    test_df: pd.DataFrame | None = None,
    association_threshold: float = 0.30,
    shift_threshold: float = 0.20,
    min_missing_rate: float = 0.01,
) -> list[LeakIssue]:
    """
    Detect columns where missingness itself appears predictive of the target,
    or where null patterns differ sharply between train and test.

    v1 signals:
    - null-indicator association with target
    - sharp train/test null-rate differences
    """
    issues: list[LeakIssue] = []

    if df.empty or target not in df.columns:
        return issues
    
    df = df.sample(10000)

    target_series = df[target]
    target_is_numeric = is_numeric_dtype(target_series)

    for col in df.columns:
        if col == target:
            continue

        series = df[col]
        missing_rate = float(series.isna().mean()) if len(series) else 0.0

        if missing_rate < min_missing_rate:
            continue

        mask = series.isna()

        if target_is_numeric:
            assoc = _null_target_association_numeric_target(mask, target_series)
        else:
            assoc = _null_target_association_categorical_target(mask, target_series)

        metric_name = assoc.get("metric_name")
        metric_value = assoc.get("metric_value")

        null_rate_diff = None
        if train_df is not None and test_df is not None and col in train_df.columns and col in test_df.columns:
            null_rate_diff = _null_rate_shift(train_df[col], test_df[col])

        severity = None
        reasons: list[str] = []
        details: dict[str, Any] = {
            "missing_rate": round(missing_rate, 6),
        }

        if metric_name is not None:
            details[metric_name] = metric_value

        if assoc.get("null_target_mean") is not None:
            details["null_target_mean"] = assoc["null_target_mean"]
        if assoc.get("non_null_target_mean") is not None:
            details["non_null_target_mean"] = assoc["non_null_target_mean"]

        if null_rate_diff is not None:
            details["train_test_missing_rate_diff"] = round(null_rate_diff, 6)

        if metric_value is not None and float(metric_value) >= association_threshold:
            severity = "high"
            reasons.append("missingness_associated_with_target")

        if null_rate_diff is not None and null_rate_diff >= shift_threshold:
            if severity != "high":
                severity = "medium"
            reasons.append("train_test_missingness_shift")

        if severity is None:
            continue

        details["reasons"] = reasons

        issues.append(
            LeakIssue(
                severity=severity,
                check="missingness_leakage",
                title="Missingness pattern may leak information",
                summary=(
                    f"Missingness in column `{col}` appears informative and may encode "
                    f"target-related or post-outcome information."
                ),
                affected_columns=[col],
                details=details,
                suggested_fix=(
                    "Investigate when and why values become missing. Missingness can leak "
                    "future workflow state, manual review status, or post-event data availability."
                ),
            )
        )

    return issues