from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
from pandas.api.types import (
    is_bool_dtype,
    is_integer_dtype,
    is_numeric_dtype,
)

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


ID_NAME_TOKENS = (
    "id",
    "uuid",
    "guid",
    "key",
    "code",
    "token",
    "hash",
    "identifier",
    "account",
    "customer",
    "user",
    "session",
    "transaction",
    "order",
    "claim",
    "policy",
    "store",
    "client",
)


def _single_column_target_signal(
    feature: pd.Series,
    target: pd.Series,
    *,
    max_categories_for_group_check: int = 5000,
) -> dict[str, Any]:
    """
    Lightweight association metrics for identifier suspicion.

    Behavior:
    - numeric target + numeric feature -> absolute Pearson correlation
    - low/moderate-cardinality feature -> target group mean spread
    """
    out: dict[str, Any] = {
        "metric_name": None,
        "metric_value": None,
    }

    try:
        if is_numeric_dtype(target) and is_numeric_dtype(feature):
            numeric_feature = pd.to_numeric(feature, errors="coerce")
            numeric_target = pd.to_numeric(target, errors="coerce")
            mask = numeric_feature.notna() & numeric_target.notna()
            if int(mask.sum()) >= 2:
                corr = numeric_feature[mask].corr(numeric_target[mask])
                out["metric_name"] = "abs_target_corr"
                out["metric_value"] = None if pd.isna(corr) else float(abs(corr))
                return out
    except Exception:
        pass

    try:
        safe_feature = feature.astype("object").where(feature.notna(), "__NaN__").astype(str)
        n_groups = int(safe_feature.nunique(dropna=False))
        if n_groups <= 1 or n_groups > max_categories_for_group_check:
            return out

        tmp = pd.DataFrame({"feature": safe_feature, "target": target})
        grouped = tmp.groupby("feature")["target"].mean()

        if len(grouped) > 1:
            spread = grouped.max() - grouped.min()
            out["metric_name"] = "target_group_mean_spread"
            out["metric_value"] = float(spread)
    except Exception:
        pass

    return out


def _is_integer_like_numeric(series: pd.Series, sample_size: int = 5000) -> bool:
    """
    Returns True when a numeric series is effectively integer-coded.
    Helps distinguish IDs/codes from continuous measurements.
    """
    try:
        s = pd.to_numeric(series, errors="coerce").dropna()
        if s.empty:
            return False

        if len(s) > sample_size:
            s = s.sample(sample_size, random_state=42)

        frac = np.abs(s.to_numpy(dtype=float) - np.round(s.to_numpy(dtype=float)))
        return bool(np.mean(frac < 1e-9) >= 0.995)
    except Exception:
        return False


def _numeric_spacing_is_id_like(series: pd.Series, sample_size: int = 10000) -> bool:
    """
    Heuristic for numeric identifiers:
    - mostly unique
    - integer-like
    - narrow gap structure (e.g. sequential IDs / codes)

    Avoids flagging continuous floats like study_hours or final_grade.
    """
    try:
        s = pd.to_numeric(series, errors="coerce").dropna()
        if len(s) < 20:
            return False

        if len(s) > sample_size:
            s = s.sample(sample_size, random_state=42)

        s = s.drop_duplicates().sort_values()
        if len(s) < 20:
            return False

        if not _is_integer_like_numeric(s):
            return False

        diffs = np.diff(s.to_numpy(dtype=float))
        diffs = diffs[diffs > 0]
        if len(diffs) == 0:
            return False

        share_small_integer_steps = np.mean(np.isin(diffs, [1.0, 2.0, 5.0, 10.0]))
        return bool(share_small_integer_steps >= 0.80)
    except Exception:
        return False


def _is_opaque_high_cardinality_string(
    inferred_type: str,
    unique_ratio: float,
    avg_len: float,
) -> bool:
    return (
        inferred_type in {"categorical", "text"}
        and unique_ratio >= 0.80
        and avg_len >= 12
    )


def _should_consider_as_identifier_candidate(
    *,
    col: str,
    info: dict[str, Any],
    suspicious_name_tokens: tuple[str, ...],
    unique_ratio_threshold: float,
    df: pd.DataFrame,
) -> tuple[bool, dict[str, bool]]:
    """
    Candidate selection with safer numeric handling.
    """
    col_lower = str(col).lower()
    unique_ratio = float(info.get("unique_ratio", 0.0))
    avg_len = float(info.get("avg_string_length", 0.0))
    inferred_type = str(info.get("inferred_type", "unknown"))
    identifier_like = bool(info.get("identifier_like", False))

    series = df[col]
    numeric = is_numeric_dtype(series) and not is_bool_dtype(series)

    name_suspicious = any(token in col_lower for token in suspicious_name_tokens)
    opaque_string_like = _is_opaque_high_cardinality_string(
        inferred_type=inferred_type,
        unique_ratio=unique_ratio,
        avg_len=avg_len,
    )

    # Important change:
    # Near-unique numeric columns are NOT identifier candidates by default.
    near_unique_string = (not numeric) and unique_ratio >= unique_ratio_threshold
    near_unique_numeric = numeric and unique_ratio >= 0.9995 and _numeric_spacing_is_id_like(series)

    should_consider = any(
        [
            identifier_like,
            name_suspicious,
            opaque_string_like,
            near_unique_string,
            near_unique_numeric,
        ]
    )

    flags = {
        "name_suspicious": name_suspicious,
        "opaque_string_like": opaque_string_like,
        "near_unique_string": near_unique_string,
        "near_unique_numeric": near_unique_numeric,
        "identifier_like": identifier_like,
        "numeric": numeric,
    }
    return should_consider, flags


def detect_identifier_leakage(
    df: pd.DataFrame,
    target: str | None = None,
    unique_ratio_threshold: float = 0.95,
    suspicious_name_tokens: tuple[str, ...] = ID_NAME_TOKENS,
) -> list[LeakIssue]:
    """
    Detect columns that behave like identifiers and may let the model memorize.

    Safer behavior than the previous version:
    - high-cardinality strings are still flagged
    - suspicious ID-like names are still flagged
    - schema-derived identifier columns are still flagged
    - continuous numeric measurements are NOT flagged just for being near-unique
    - numeric columns need extra evidence (e.g. suspicious name, schema flag,
      or integer/sequential ID-like spacing)
    """
    issues: list[LeakIssue] = []

    if df.empty:
        return issues

    schema = infer_schema(df)
    column_info = schema["columns"]

    candidate_columns: list[str] = []
    candidate_flags: dict[str, dict[str, bool]] = {}

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

        info = column_info.get(col, {})
        should_consider, flags = _should_consider_as_identifier_candidate(
            col=col,
            info=info,
            suspicious_name_tokens=suspicious_name_tokens,
            unique_ratio_threshold=unique_ratio_threshold,
            df=df,
        )

        if should_consider:
            candidate_columns.append(col)
            candidate_flags[col] = flags

    candidate_columns = candidate_columns[:max_target_signal_columns]

    for col in candidate_columns:
        info = column_info.get(col, {})
        flags = candidate_flags.get(col, {})

        unique_ratio = float(info.get("unique_ratio", 0.0))
        nunique = int(info.get("nunique", 0))
        avg_len = float(info.get("avg_string_length", 0.0))
        inferred_type = str(info.get("inferred_type", "unknown"))
        identifier_like = bool(info.get("identifier_like", False))

        series = df[col]
        numeric = bool(flags.get("numeric", False))
        name_suspicious = bool(flags.get("name_suspicious", False))
        opaque_string_like = bool(flags.get("opaque_string_like", False))
        near_unique_string = bool(flags.get("near_unique_string", False))
        near_unique_numeric = bool(flags.get("near_unique_numeric", False))

        reasons: list[str] = []
        severity = "medium"

        if name_suspicious:
            reasons.append("suspicious_name")
        if near_unique_string or near_unique_numeric:
            reasons.append("near_unique_values")
        if opaque_string_like:
            reasons.append("opaque_high_cardinality_strings")
        if identifier_like:
            reasons.append("schema_identifier_like")

        # Severity rules
        if name_suspicious and (identifier_like or near_unique_numeric or near_unique_string):
            severity = "high"
        elif identifier_like and unique_ratio >= 0.995:
            severity = "high"
        elif opaque_string_like and unique_ratio >= 0.98:
            severity = "high"
        elif near_unique_numeric and unique_ratio >= 0.9999:
            severity = "high"

        details: dict[str, Any] = {
            "nunique": nunique,
            "unique_ratio": round(unique_ratio, 6),
            "avg_string_length": round(avg_len, 3),
            "inferred_type": inferred_type,
            "reasons": list(reasons),
        }

        if numeric:
            details["integer_like"] = _is_integer_like_numeric(series)
            details["sequential_id_like_spacing"] = _numeric_spacing_is_id_like(series)

        if target is not None and target in df.columns:
            signal = _single_column_target_signal(series, df[target])
            metric_name = signal.get("metric_name")
            metric_value = signal.get("metric_value")

            if metric_name is not None:
                details[metric_name] = None if metric_value is None else round(float(metric_value), 6)

                # Important: target association should only escalate a column that
                # already looks identifier-like for structural reasons.
                structural_identifier_evidence = any(
                    [
                        name_suspicious,
                        identifier_like,
                        opaque_string_like,
                        near_unique_numeric,
                        near_unique_string,
                    ]
                )

                if structural_identifier_evidence:
                    if metric_name == "abs_target_corr" and metric_value is not None and metric_value >= 0.30:
                        severity = "high"
                        reasons.append("predictive_identifier_signal")

                    if (
                        metric_name == "target_group_mean_spread"
                        and metric_value is not None
                        and metric_value > 0
                    ):
                        reasons.append("target_group_signal")

        details["reasons"] = sorted(set(reasons))

        issues.append(
            LeakIssue(
                severity=severity,
                check="identifier_leakage",
                title="Identifier-like feature detected",
                summary=(
                    f"Column `{col}` behaves like an identifier and may allow memorization "
                    f"instead of genuine generalization."
                ),
                affected_columns=[col],
                details=details,
                suggested_fix=(
                    "Drop this feature from training, or validate using entity/group-based splits "
                    "if the identifier represents a real unit such as customer, store, or account."
                ),
            )
        )

    return issues