from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split

from featureleak.report import LeakIssue
from featureleak.utils.encoding import lightweight_encode_frame


def _prepare_adversarial_frame(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    target: str | None = None,
    time_col: str | None = None,
    exclude_columns: list[str] | None = None,
) -> tuple[pd.DataFrame, pd.Series, list[str]]:
    exclude = set(exclude_columns or [])
    if target is not None:
        exclude.add(target)
    if time_col is not None:
        exclude.add(time_col)

    shared_cols = [c for c in train_df.columns if c in test_df.columns and c not in exclude]

    train_x = train_df[shared_cols].copy()
    test_x = test_df[shared_cols].copy()

    combined = pd.concat([train_x, test_x], axis=0, ignore_index=True)
    y = pd.Series(
        np.concatenate(
            [
                np.zeros(len(train_x), dtype=int),
                np.ones(len(test_x), dtype=int),
            ]
        ),
        name="is_test",
    )

    return combined, y, shared_cols


def _sample_if_needed(
    X: pd.DataFrame,
    y: pd.Series,
    sample_size: int = 50_000,
    random_state: int = 42,
) -> tuple[pd.DataFrame, pd.Series]:
    if len(X) <= sample_size:
        return X, y

    sampled_idx = (
        pd.Series(np.arange(len(X)))
        .sample(n=sample_size, random_state=random_state, replace=False)
        .sort_values()
        .to_numpy()
    )
    return X.iloc[sampled_idx].reset_index(drop=True), y.iloc[sampled_idx].reset_index(drop=True)


def detect_adversarial_shift(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    target: str | None = None,
    time_col: str | None = None,
    exclude_columns: list[str] | None = None,
    sample_size: int = 50_000,
    auc_warning_threshold: float = 0.70,
    auc_high_threshold: float = 0.85,
    random_state: int = 42,
) -> list[LeakIssue]:
    """
    Train a model to distinguish train from test.

    High AUC means the split is separable, suggesting:
    - dataset shift
    - broken preprocessing alignment
    - temporal drift
    - leakage via split construction
    """
    issues: list[LeakIssue] = []

    if train_df.empty or test_df.empty:
        return issues

    X_raw, y, shared_cols = _prepare_adversarial_frame(
        train_df=train_df,
        test_df=test_df,
        target=target,
        time_col=time_col,
        exclude_columns=exclude_columns,
    )

    if X_raw.empty or len(shared_cols) == 0:
        return issues

    X_raw, y = _sample_if_needed(
        X=X_raw,
        y=y,
        sample_size=sample_size,
        random_state=random_state,
    )

    X = lightweight_encode_frame(X_raw)

    if X.empty:
        return issues

    # Handle rare edge case where only one class survives, though unlikely.
    if y.nunique() < 2:
        return issues

    X_train, X_valid, y_train, y_valid = train_test_split(
        X,
        y,
        test_size=0.25,
        random_state=random_state,
        stratify=y,
    )

    model = RandomForestClassifier(
        n_estimators=200,
        max_depth=8,
        min_samples_leaf=20,
        n_jobs=-1,
        random_state=random_state,
    )
    model.fit(X_train, y_train)

    proba = model.predict_proba(X_valid)[:, 1]
    auc = float(roc_auc_score(y_valid, proba))

    importances = getattr(model, "feature_importances_", None)
    top_features: list[str] = []
    if importances is not None:
        ranked = (
            pd.Series(importances, index=X.columns)
            .sort_values(ascending=False)
            .head(5)
        )
        top_features = ranked.index.tolist()

    if auc < auc_warning_threshold:
        return issues

    severity = "high" if auc >= auc_high_threshold else "medium"

    summary = (
        f"Train/test datasets are distinguishable (AUC={auc:.3f}). "
        f"This suggests leakage, dataset shift, or inconsistent preprocessing."
    )

    issues.append(
        LeakIssue(
            severity=severity,
            check="adversarial_validation",
            title="Train/test split is highly separable",
            summary=summary,
            affected_columns=top_features,
            details={
                "auc": round(auc, 6),
                "top_features": top_features,
                "n_shared_features": len(shared_cols),
                "sample_size_used": int(len(X)),
                "excluded_columns": sorted(list(set(exclude_columns or []))),
            },
            suggested_fix=(
                "Inspect the top distinguishing features between train and test. "
                "They may reveal temporal leakage, preprocessing differences, "
                "entity leakage, or dataset shift."
            ),
        )
    )

    return issues