from __future__ import annotations

import numpy as np
import pandas as pd

from sklearn.decomposition import PCA
from sklearn.impute import SimpleImputer
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import StratifiedKFold, cross_val_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from lightgbm import LGBMClassifier

from featureleak.report import WarningItem


def _numeric_only(df: pd.DataFrame) -> pd.DataFrame:
    if df is None:
        return pd.DataFrame()

    out = df.copy()

    datetime_cols = out.select_dtypes(include=["datetime64[ns]", "datetime64"]).columns
    for col in datetime_cols:
        out[col] = out[col].astype("int64")

    for col in out.columns:
        if out[col].dtype == "object":
            parsed = pd.to_datetime(out[col], errors="coerce")
            if parsed.notna().mean() > 0.8:
                out[col] = parsed.astype("int64")

    out = pd.get_dummies(out, dummy_na=True)

    bool_cols = out.select_dtypes(include=["bool"]).columns
    for col in bool_cols:
        out[col] = out[col].astype(int)

    return out.select_dtypes(include=[np.number]).copy()


def _align_columns(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    train_aligned, test_aligned = train_df.align(test_df, join="outer", axis=1, fill_value=np.nan)
    return train_aligned, test_aligned


def _mean_abs_diff(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.nanmean(np.abs(a - b)))


def _safe_auc(
    X: pd.DataFrame,
    y: pd.Series,
    random_state: int,
) -> float | None:
    if X.empty or y.nunique() < 2:
        return None

    model = LGBMClassifier(
        n_estimators=100,
        learning_rate=0.05,
        num_leaves=31,
        subsample=0.8,
        colsample_bytree=0.8,
        random_state=random_state,
        n_jobs=-1,
        verbose=-1,
    )

    cv = StratifiedKFold(n_splits=3, shuffle=True, random_state=random_state)

    try:
        scores = cross_val_score(
            model,
            X,
            y,
            cv=cv,
            scoring="roc_auc",
            error_score="raise",
        )
        return float(np.mean(scores))
    except Exception:
        return None


def _check_suspicious_global_scaling(
    X_train_processed: pd.DataFrame | None,
    X_test_processed: pd.DataFrame | None,
    near_zero_tol: float = 0.05,
    near_one_tol: float = 0.05,
    min_flagged_features: int = 5,
) -> list[WarningItem]:
    warnings: list[WarningItem] = []

    if X_train_processed is None or X_test_processed is None:
        return warnings

    train_num = _numeric_only(X_train_processed)
    test_num = _numeric_only(X_test_processed)

    if train_num.empty or test_num.empty:
        return warnings

    train_num, test_num = _align_columns(train_num, test_num)

    flagged: list[str] = []
    details: dict[str, dict[str, float]] = {}

    for col in train_num.columns:
        train_col = pd.to_numeric(train_num[col], errors="coerce")
        test_col = pd.to_numeric(test_num[col], errors="coerce")

        if train_col.notna().sum() < 5 or test_col.notna().sum() < 5:
            continue

        train_mean = float(train_col.mean())
        test_mean = float(test_col.mean())
        train_std = float(train_col.std(ddof=0))
        test_std = float(test_col.std(ddof=0))

        looks_globally_standardized = (
            abs(train_mean) <= near_zero_tol
            and abs(test_mean) <= near_zero_tol
            and abs(train_std - 1.0) <= near_one_tol
            and abs(test_std - 1.0) <= near_one_tol
        )

        if looks_globally_standardized:
            flagged.append(col)
            details[col] = {
                "train_mean": train_mean,
                "test_mean": test_mean,
                "train_std": train_std,
                "test_std": test_std,
            }

    if len(flagged) >= min_flagged_features:
        warnings.append(
            WarningItem(
                rule="preprocessing_contamination_scaling_pattern",
                severity="medium",
                message=(
                    f"{len(flagged)} processed numeric features appear standardized in both "
                    f"train and test. This may indicate scaling fit on all rows before splitting."
                ),
                columns=flagged[:10],
                details={
                    "n_flagged_features": int(len(flagged)),
                    "examples": dict(list(details.items())[:5]),
                },
                suggested_fix=(
                    "Fit StandardScaler or similar transformers on training data only, "
                    "then transform the test set using the fitted training transformer."
                ),
            )
        )

    return warnings


def _check_transform_reconstruction(
    X_train_raw: pd.DataFrame | None,
    X_test_raw: pd.DataFrame | None,
    X_test_processed: pd.DataFrame | None,
    max_features_for_pca: int = 50,
) -> list[WarningItem]:
    warnings: list[WarningItem] = []

    if X_train_raw is None or X_test_raw is None or X_test_processed is None:
        return warnings

    train_raw = _numeric_only(X_train_raw)
    test_raw = _numeric_only(X_test_raw)
    test_proc = _numeric_only(X_test_processed)

    if train_raw.empty or test_raw.empty or test_proc.empty:
        return warnings

    train_raw, test_raw = _align_columns(train_raw, test_raw)
    _, test_proc = _align_columns(train_raw, test_proc)

    if train_raw.shape[1] == 0:
        return warnings

    transform_candidates: list[tuple[str, object]] = [
        ("standard_scaler", StandardScaler()),
        ("minmax_scaler", MinMaxScaler()),
        ("median_imputer", SimpleImputer(strategy="median")),
    ]

    if train_raw.shape[1] >= 2:
        n_comp = min(5, train_raw.shape[1], max_features_for_pca)
        if n_comp >= 2:
            transform_candidates.append(
                (
                    "median_imputer_plus_pca",
                    Pipeline(
                        steps=[
                            ("imputer", SimpleImputer(strategy="median")),
                            ("pca", PCA(n_components=n_comp, random_state=42)),
                        ]
                    ),
                )
            )

    results: list[dict[str, float | str]] = []

    for name, transformer in transform_candidates:
        try:
            train_fit = transformer.fit(train_raw)
            test_trainfit = train_fit.transform(test_raw)

            all_rows = pd.concat([train_raw, test_raw], axis=0, ignore_index=True)
            all_fit = transformer.fit(all_rows)
            test_allfit = all_fit.transform(test_raw)

            # Compare only when dimensions match the provided processed test
            proc_values = test_proc.to_numpy()

            if test_trainfit.shape != proc_values.shape or test_allfit.shape != proc_values.shape:
                continue

            dist_trainfit = _mean_abs_diff(proc_values, test_trainfit)
            dist_allfit = _mean_abs_diff(proc_values, test_allfit)

            results.append(
                {
                    "transform": name,
                    "distance_trainfit": float(dist_trainfit),
                    "distance_allfit": float(dist_allfit),
                    "ratio_allfit_to_trainfit": float(
                        dist_allfit / max(dist_trainfit, 1e-12)
                    ),
                }
            )
        except Exception:
            continue

    suspicious = [
        r for r in results
        if r["distance_allfit"] < r["distance_trainfit"] * 0.5
    ]

    if suspicious:
        best = sorted(suspicious, key=lambda x: x["ratio_allfit_to_trainfit"])[0]

        warnings.append(
            WarningItem(
                rule="preprocessing_contamination_reconstruction",
                severity="high",
                message=(
                    "Processed test features are materially closer to a transformer fit on "
                    "train+test than to one fit on train only. This strongly suggests "
                    "pre-split preprocessing contamination."
                ),
                columns=[],
                details={
                    "best_matching_transform": best["transform"],
                    "distance_trainfit": best["distance_trainfit"],
                    "distance_allfit": best["distance_allfit"],
                    "ratio_allfit_to_trainfit": best["ratio_allfit_to_trainfit"],
                    "tested_transforms": results,
                },
                suggested_fix=(
                    "Move preprocessing inside a leak-safe sklearn Pipeline and fit it "
                    "only on training folds or the training split."
                ),
            )
        )

    return warnings


def _check_validation_inflation(
    X_train_raw: pd.DataFrame | None,
    X_train_processed: pd.DataFrame | None,
    y_train: pd.Series | None,
    random_state: int = 42,
    inflation_threshold: float = 0.03,
) -> list[WarningItem]:
    warnings: list[WarningItem] = []

    if X_train_raw is None or X_train_processed is None or y_train is None:
        return warnings

    raw_num = _numeric_only(X_train_raw)
    proc_num = _numeric_only(X_train_processed)

    if raw_num.empty or proc_num.empty:
        return warnings

    y = pd.Series(y_train).reset_index(drop=True)
    raw_num = raw_num.reset_index(drop=True)
    proc_num = proc_num.reset_index(drop=True)

    if len(raw_num) != len(y) or len(proc_num) != len(y):
        return warnings

    proper_pipeline = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="median")),
            ("scaler", StandardScaler()),
            (
                "model",
                LGBMClassifier(
                    n_estimators=100,
                    learning_rate=0.05,
                    num_leaves=31,
                    subsample=0.8,
                    colsample_bytree=0.8,
                    random_state=random_state,
                    n_jobs=-1,
                    verbose=-1,
                ),
            ),
        ]
    )

    try:
        cv = StratifiedKFold(n_splits=3, shuffle=True, random_state=random_state)

        proper_scores = cross_val_score(
            proper_pipeline,
            raw_num,
            y,
            cv=cv,
            scoring="roc_auc",
            error_score="raise",
        )
        proper_auc = float(np.mean(proper_scores))
    except Exception:
        proper_auc = None

    contaminated_auc = _safe_auc(proc_num, y, random_state=random_state)

    if proper_auc is None or contaminated_auc is None:
        return warnings

    inflation = contaminated_auc - proper_auc

    if inflation > inflation_threshold:
        warnings.append(
            WarningItem(
                rule="preprocessing_contamination_validation_inflation",
                severity="medium",
                message=(
                    f"Validation performance on processed training data is inflated "
                    f"relative to a leak-safe pipeline (ΔAUC={inflation:.3f})."
                ),
                columns=[],
                details={
                    "processed_auc": contaminated_auc,
                    "proper_pipeline_auc": proper_auc,
                    "inflation": inflation,
                },
                suggested_fix=(
                    "Evaluate models using a single end-to-end Pipeline inside cross-validation "
                    "so transformers are refit within each fold."
                ),
            )
        )

    return warnings


def detect_preprocessing_contamination(
    train_df: pd.DataFrame | None = None,
    test_df: pd.DataFrame | None = None,
    target: str | None = None,
    X_train_raw: pd.DataFrame | None = None,
    X_test_raw: pd.DataFrame | None = None,
    X_train_processed: pd.DataFrame | None = None,
    X_test_processed: pd.DataFrame | None = None,
    y_train: pd.Series | None = None,
) -> list[WarningItem]:
    warnings: list[WarningItem] = []

    # Fallbacks for easy usage
    if X_train_raw is None and train_df is not None:
        X_train_raw = train_df.drop(columns=[target], errors="ignore")

    if X_test_raw is None and test_df is not None:
        X_test_raw = test_df.drop(columns=[target], errors="ignore")

    if y_train is None and train_df is not None and target is not None and target in train_df.columns:
        y_train = train_df[target]

    warnings.extend(
        _check_suspicious_global_scaling(
            X_train_processed=X_train_processed,
            X_test_processed=X_test_processed,
        )
    )

    warnings.extend(
        _check_transform_reconstruction(
            X_train_raw=X_train_raw,
            X_test_raw=X_test_raw,
            X_test_processed=X_test_processed,
        )
    )

    warnings.extend(
        _check_validation_inflation(
            X_train_raw=X_train_raw,
            X_train_processed=X_train_processed,
            y_train=y_train,
        )
    )

    return warnings