from __future__ import annotations

from typing import Any

import pandas as pd
from pandas.api.types import (
    is_datetime64_any_dtype,
    is_numeric_dtype,
)
from featureleak.utils.datetime import safe_to_datetime

from featureleak.report import LeakIssue


def _safe_dtype_name(series: pd.Series) -> str:
    try:
        return str(series.dtype)
    except Exception:
        return "unknown"


def _missing_rate(series: pd.Series) -> float:
    try:
        return float(series.isna().mean())
    except Exception:
        return 0.0


def _safe_numeric_summary(series: pd.Series) -> dict[str, float | None]:
    try:
        numeric = pd.to_numeric(series, errors="coerce")
        if numeric.notna().sum() == 0:
            return {
                "mean": None,
                "std": None,
                "min": None,
                "max": None,
            }
        return {
            "mean": None if pd.isna(numeric.mean()) else float(numeric.mean()),
            "std": None if pd.isna(numeric.std()) else float(numeric.std()),
            "min": None if pd.isna(numeric.min()) else float(numeric.min()),
            "max": None if pd.isna(numeric.max()) else float(numeric.max()),
        }
    except Exception:
        return {
            "mean": None,
            "std": None,
            "min": None,
            "max": None,
        }


def _safe_category_overlap(train_series: pd.Series, test_series: pd.Series) -> dict[str, Any]:
    try:
        train_vals = set(
            train_series.astype("object").where(train_series.notna(), "__NaN__").astype(str).unique().tolist()
        )
        test_vals = set(
            test_series.astype("object").where(test_series.notna(), "__NaN__").astype(str).unique().tolist()
        )

        overlap = train_vals & test_vals
        union = train_vals | test_vals

        return {
            "n_train_categories": len(train_vals),
            "n_test_categories": len(test_vals),
            "n_overlap_categories": len(overlap),
            "jaccard_overlap": len(overlap) / max(len(union), 1),
            "train_only_categories": sorted(list(train_vals - test_vals))[:10],
            "test_only_categories": sorted(list(test_vals - train_vals))[:10],
        }
    except Exception:
        return {
            "n_train_categories": 0,
            "n_test_categories": 0,
            "n_overlap_categories": 0,
            "jaccard_overlap": 0.0,
            "train_only_categories": [],
            "test_only_categories": [],
        }


def _datetime_parse_rate(series: pd.Series) -> float:
    """
    Useful for object/string columns that may contain dates in only one split.
    """
    try:
        parsed = safe_to_datetime(series)
        return float(parsed.notna().mean())
    except Exception:
        return 0.0


def detect_preprocessing_inconsistency(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    target: str | None = None,
    missing_rate_diff_threshold: float = 0.20,
    numeric_mean_shift_threshold: float = 3.0,
    category_overlap_threshold: float = 0.20,
    datetime_parse_gap_threshold: float = 0.50,
) -> list[LeakIssue]:
    """
    Detect preprocessing inconsistencies between train and test.

    v1 signals:
    - missing columns in one split
    - dtype mismatch
    - numeric scale/mean mismatch
    - large missingness differences
    - severe category-set mismatch
    - date parsing mismatch
    """
    issues: list[LeakIssue] = []

    train_cols = set(train_df.columns)
    test_cols = set(test_df.columns)

    if target is not None:
        train_cols.discard(target)
        test_cols.discard(target)

    train_only = sorted(list(train_cols - test_cols))
    test_only = sorted(list(test_cols - train_cols))

    if train_only or test_only:
        issues.append(
            LeakIssue(
                severity="high",
                check="preprocessing_inconsistency",
                title="Columns differ between train and test",
                summary=(
                    "Train and test do not share the same feature columns, which suggests "
                    "inconsistent preprocessing or schema drift."
                ),
                affected_columns=train_only + test_only,
                details={
                    "train_only_columns": train_only,
                    "test_only_columns": test_only,
                    "n_train_only_columns": len(train_only),
                    "n_test_only_columns": len(test_only),
                },
                suggested_fix=(
                    "Ensure train and test are generated by the same preprocessing pipeline "
                    "and aligned to the same feature schema before model training."
                ),
            )
        )

    shared_cols = sorted(list(train_cols & test_cols))

    for col in shared_cols:
        train_s = train_df[col]
        test_s = test_df[col]

        train_dtype = _safe_dtype_name(train_s)
        test_dtype = _safe_dtype_name(test_s)

        reasons: list[str] = []
        details: dict[str, Any] = {
            "train_dtype": train_dtype,
            "test_dtype": test_dtype,
        }
        severity: str | None = None

        # --------------------------------------------------------------
        # Dtype mismatch
        # --------------------------------------------------------------
        if train_dtype != test_dtype:
            reasons.append("dtype_mismatch")
            severity = "medium"

        # --------------------------------------------------------------
        # Missingness mismatch
        # --------------------------------------------------------------
        train_missing = _missing_rate(train_s)
        test_missing = _missing_rate(test_s)
        missing_diff = abs(train_missing - test_missing)

        details["train_missing_rate"] = round(train_missing, 6)
        details["test_missing_rate"] = round(test_missing, 6)
        details["missing_rate_diff"] = round(missing_diff, 6)

        if missing_diff >= missing_rate_diff_threshold:
            reasons.append("missingness_shift")
            severity = "high" if missing_diff >= 0.40 else max(severity or "low", "medium", key=lambda x: ["low","medium","high"].index(x))

        # --------------------------------------------------------------
        # Numeric distribution mismatch
        # --------------------------------------------------------------
        if is_numeric_dtype(train_s) and is_numeric_dtype(test_s):
            train_stats = _safe_numeric_summary(train_s)
            test_stats = _safe_numeric_summary(test_s)

            for k, v in train_stats.items():
                details[f"train_{k}"] = v
            for k, v in test_stats.items():
                details[f"test_{k}"] = v

            train_std = train_stats.get("std")
            train_mean = train_stats.get("mean")
            test_mean = test_stats.get("mean")

            if (
                train_std is not None
                and train_std > 0
                and train_mean is not None
                and test_mean is not None
            ):
                normalized_mean_gap = abs(test_mean - train_mean) / train_std
                details["normalized_mean_gap"] = float(normalized_mean_gap)

                if normalized_mean_gap >= numeric_mean_shift_threshold:
                    reasons.append("numeric_distribution_shift")
                    severity = "medium" if severity is None else severity

        # --------------------------------------------------------------
        # Datetime mismatch / parse mismatch
        # --------------------------------------------------------------
        elif is_datetime64_any_dtype(train_s) != is_datetime64_any_dtype(test_s):
            reasons.append("datetime_type_mismatch")
            severity = "medium" if severity is None else severity

        else:
            # Categorical / object-style checks
            overlap = _safe_category_overlap(train_s, test_s)
            details.update(overlap)

            jaccard = float(overlap["jaccard_overlap"])
            if jaccard <= category_overlap_threshold:
                reasons.append("category_set_mismatch")
                severity = "medium" if severity is None else severity

            # Object/string date parsing mismatch
            train_parse_rate = _datetime_parse_rate(train_s)
            test_parse_rate = _datetime_parse_rate(test_s)
            parse_gap = abs(train_parse_rate - test_parse_rate)

            details["train_datetime_parse_rate"] = round(train_parse_rate, 6)
            details["test_datetime_parse_rate"] = round(test_parse_rate, 6)
            details["datetime_parse_rate_gap"] = round(parse_gap, 6)

            if parse_gap >= datetime_parse_gap_threshold:
                reasons.append("date_parsing_inconsistency")
                severity = "medium" if severity is None else severity

        if severity is None:
            continue

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

        issues.append(
            LeakIssue(
                severity=severity,
                check="preprocessing_inconsistency",
                title="Preprocessing mismatch detected",
                summary=(
                    f"Column `{col}` looks inconsistently processed between train and test."
                ),
                affected_columns=[col],
                details=details,
                suggested_fix=(
                    "Fit preprocessing on training data only and apply the exact same fitted "
                    "pipeline to test data. Check dtype coercion, missing-value handling, "
                    "category mapping, and date parsing."
                ),
            )
        )

    return issues