from __future__ import annotations

import time
import warnings
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from typing import Any

import pandas as pd
from tqdm import tqdm

from featureleak.report import LeakIssue, LeakReport, WarningItem
from featureleak.rules.adversarial_validation import detect_adversarial_shift
from featureleak.rules.aggregation_leakage import detect_aggregation_leakage
from featureleak.rules.contamination import detect_contamination
from featureleak.rules.duplicates import detect_duplicate_rows
from featureleak.rules.entity_leakage import detect_entity_leakage
from featureleak.rules.identifier_leakage import detect_identifier_leakage
from featureleak.rules.missingness_leakage import detect_missingness_leakage
from featureleak.rules.preprocessing_inconsistency import detect_preprocessing_inconsistency
from featureleak.rules.target_leakage import detect_target_leakage
from featureleak.rules.time_leakage import detect_time_leakage


SEVERITY_ORDER: dict[str, int] = {
    "critical": 4,
    "high": 3,
    "medium": 2,
    "low": 1,
    "info": 0,
}

CHECK_ORDER: dict[str, int] = {
    "target_leakage": 100,
    "temporal_leakage": 90,
    "train_test_contamination": 80,
    "duplicate_leakage": 70,
    "entity_leakage": 60,
    "aggregation_leakage": 50,
    "identifier_leakage": 40,
    "missingness_leakage": 30,
    "adversarial_validation": 20,
    "preprocessing_inconsistency": 10,
    "analysis_step_error": 0,
}


@dataclass
class ScannerConfig:
    target_corr_threshold: float = 0.98
    overlap_threshold: float = 0.0
    sample_size: int | None = None


class LeakScanner:
    def __init__(
        self,
        target_corr_threshold: float = 0.98,
        overlap_threshold: float = 0.0,
        sample_size: int | None = None,
    ) -> None:
        self.config = ScannerConfig(
            target_corr_threshold=target_corr_threshold,
            overlap_threshold=overlap_threshold,
            sample_size=sample_size,
        )

    def _in_notebook(self) -> bool:
        try:
            from IPython import get_ipython

            shell = get_ipython()
            if shell is None:
                return False
            return shell.__class__.__name__ in {"ZMQInteractiveShell", "Shell"}
        except Exception:
            return False

    def _build_steps(
        self,
        df: pd.DataFrame,
        target: str,
        time_col: str | None,
        train_df: pd.DataFrame | None,
        test_df: pd.DataFrame | None,
        entity_cols: list[str] | None,
    ) -> list[tuple[str, str, Callable[[], Any]]]:
        steps: list[tuple[str, str, Callable[[], Any]]] = []

        steps.append(
            (
                "target_leakage",
                "Checking target leakage",
                lambda: detect_target_leakage(
                    df=df,
                    target=target,
                    corr_threshold=self.config.target_corr_threshold,
                ),
            )
        )

        steps.append(
            (
                "duplicate_leakage",
                "Checking duplicate and near-duplicate rows",
                lambda: detect_duplicate_rows(
                    df=df,
                    target=target,
                    sample_size=self.config.sample_size,
                ),
            )
        )

        steps.append(
            (
                "identifier_leakage",
                "Checking identifier leakage",
                lambda: detect_identifier_leakage(
                    df=df,
                    target=target,
                    sample_size=self.config.sample_size,
                ),
            )
        )

        steps.append(
            (
                "aggregation_leakage",
                "Checking target aggregation leakage",
                lambda: detect_aggregation_leakage(
                    df=df,
                    target=target,
                ),
            )
        )

        steps.append(
            (
                "missingness_leakage",
                "Checking missingness leakage",
                lambda: detect_missingness_leakage(
                    df=df,
                    target=target,
                    train_df=train_df,
                    test_df=test_df,
                ),
            )
        )

        if train_df is not None and test_df is not None:
            steps.append(
                (
                    "adversarial_validation",
                    "Running adversarial validation",
                    lambda: detect_adversarial_shift(
                        train_df=train_df,
                        test_df=test_df,
                        target=target,
                        time_col=time_col,
                        exclude_columns=[time_col] if time_col is not None else None,
                    ),
                )
            )

            steps.append(
                (
                    "preprocessing_inconsistency",
                    "Checking preprocessing inconsistency",
                    lambda: detect_preprocessing_inconsistency(
                        train_df=train_df,
                        test_df=test_df,
                        target=target,
                    ),
                )
            )

        if time_col is not None:
            steps.append(
                (
                    "temporal_leakage",
                    "Checking temporal leakage",
                    lambda: detect_time_leakage(
                        df=df,
                        target=target,
                        time_col=time_col,
                        train_df=train_df,
                        test_df=test_df,
                    ),
                )
            )

        if train_df is not None and test_df is not None:
            steps.append(
                (
                    "train_test_contamination",
                    "Checking train/test contamination",
                    lambda: detect_contamination(
                        train_df=train_df,
                        test_df=test_df,
                        target=target,
                        overlap_threshold=self.config.overlap_threshold,
                        key_columns=entity_cols,
                    ),
                )
            )

            steps.append(
                (
                    "entity_leakage",
                    "Checking entity/group leakage",
                    lambda: detect_entity_leakage(
                        train_df=train_df,
                        test_df=test_df,
                        entity_cols=entity_cols,
                        target=target,
                    ),
                )
            )

        return steps

    def _coerce_issue(self, raw_item: Any, fallback_check: str) -> LeakIssue:
        if isinstance(raw_item, LeakIssue):
            return raw_item

        if isinstance(raw_item, WarningItem):
            return LeakIssue(
                severity=raw_item.severity,
                check=getattr(raw_item, "check", None) or getattr(raw_item, "rule", None) or fallback_check,
                title=getattr(raw_item, "title", "") or getattr(raw_item, "rule", fallback_check).replace("_", " ").title(),
                summary=getattr(raw_item, "summary", "") or getattr(raw_item, "message", ""),
                affected_columns=list(
                    getattr(raw_item, "affected_columns", None)
                    or getattr(raw_item, "columns", None)
                    or []
                ),
                details=getattr(raw_item, "details", {}) or {},
                suggested_fix=getattr(raw_item, "suggested_fix", "") or "",
            )

        if isinstance(raw_item, dict):
            return LeakIssue(
                severity=str(raw_item.get("severity", "low")),
                check=str(raw_item.get("check") or raw_item.get("rule") or fallback_check),
                title=str(
                    raw_item.get("title")
                    or str(raw_item.get("check") or raw_item.get("rule") or fallback_check).replace("_", " ").title()
                ),
                summary=str(raw_item.get("summary") or raw_item.get("message") or ""),
                affected_columns=list(raw_item.get("affected_columns") or raw_item.get("columns") or []),
                details=dict(raw_item.get("details") or {}),
                suggested_fix=str(raw_item.get("suggested_fix") or ""),
            )

        raise TypeError(
            f"Rule output must be LeakIssue, WarningItem, or dict; got {type(raw_item).__name__}."
        )

    def _coerce_step_output(
        self,
        raw_output: Any,
        fallback_check: str,
    ) -> list[LeakIssue]:
        if raw_output is None:
            return []

        if isinstance(raw_output, (LeakIssue, WarningItem, dict)):
            return [self._coerce_issue(raw_output, fallback_check)]

        if isinstance(raw_output, Iterable) and not isinstance(raw_output, (str, bytes)):
            return [self._coerce_issue(item, fallback_check) for item in raw_output]

        raise TypeError(
            f"Rule output must be None, a single issue, or an iterable of issues; got {type(raw_output).__name__}."
        )

    def _add_step_error(
        self,
        report: LeakReport,
        check: str,
        description: str,
        exc: Exception,
    ) -> None:
        report.add_issue(
            LeakIssue(
                severity="low",
                check="analysis_step_error",
                title="Analysis step failed",
                summary=f"{description} failed: {type(exc).__name__}: {exc}",
                affected_columns=[],
                details={
                    "step": description,
                    "failed_check": check,
                    "error_type": type(exc).__name__,
                    "error_message": str(exc),
                },
                suggested_fix=(
                    "Check input data types, required columns, and optional dependencies "
                    "for this analysis step."
                ),
            )
        )

    def _issue_sort_key(self, issue: LeakIssue) -> tuple[int, int, str, str]:
        severity_rank = SEVERITY_ORDER.get(str(issue.severity).lower(), -1)
        check_rank = CHECK_ORDER.get(str(issue.check), -1)
        first_column = issue.affected_columns[0] if issue.affected_columns else ""
        return (
            severity_rank,
            check_rank,
            first_column,
            issue.title,
        )

    def _sort_report_issues(self, report: LeakReport) -> None:
        report.issues = sorted(
            report.issues,
            key=self._issue_sort_key,
            reverse=True,
        )

    def _finalise_risk_buckets(self, report: LeakReport) -> None:
        high: set[str] = set()
        medium: set[str] = set()
        low: set[str] = set()

        for issue in report.issues:
            sev = issue.severity.lower()
            if sev in {"critical", "high"}:
                high.update(issue.affected_columns)
            elif sev == "medium":
                medium.update(issue.affected_columns)
            elif sev == "low":
                low.update(issue.affected_columns)

        report.high_risk_features = sorted(high)
        report.medium_risk_features = sorted(medium - high)
        report.low_risk_features = sorted(low - high - medium)

    def _print_terminal_summary(
        self,
        report: LeakReport,
        step_results: list[dict[str, object]],
    ) -> None:
        print("\nFeatureLeak Scan")
        print("─" * 56)

        for result in step_results:
            status = result["status"]
            name = result["name"]
            elapsed = result["elapsed"]
            n_issues = result["n_issues"]

            icon = {"ok": "✔", "warning": "⚠", "error": "✖"}.get(status, "•")
            print(f"{icon} {name} ({elapsed:.2f}s, issues={n_issues})")

        print("─" * 56)
        print(f"Risk score: {report.score()}/100 ({report.risk_level().title()})")
        print(f"Total issues: {len(report.issues)}")
        print(f"High risk features: {len(report.high_risk_features)}")
        print(f"Medium risk features: {len(report.medium_risk_features)}")
        print(f"Low risk features: {len(report.low_risk_features)}")

        if report.issues:
            print("\nTop issues:")
            for issue in report.issues[:5]:
                print(f"- [{issue.severity.upper()}] {issue.check}: {issue.summary}")

    def analyse(
        self,
        df: pd.DataFrame,
        target: str,
        time_col: str | None = None,
        train_df: pd.DataFrame | None = None,
        test_df: pd.DataFrame | None = None,
        entity_cols: list[str] | None = None,
        show_progress: bool = True,
        show_summary: bool = True,
    ) -> LeakReport:
        if target not in df.columns:
            raise ValueError(f"Target column '{target}' not found in dataframe.")

        if time_col is not None and time_col not in df.columns:
            raise ValueError(f"Time column '{time_col}' not found in dataframe.")

        if entity_cols is not None:
            missing_entity_cols = [col for col in entity_cols if col not in df.columns]
            if missing_entity_cols:
                raise ValueError(f"Entity columns not found in dataframe: {missing_entity_cols}")

        # Sample dataframes if sample_size is set
        df_sample = df.sample(self.config.sample_size) if self.config.sample_size and self.config.sample_size < len(df) else df
        train_sample = train_df.sample(self.config.sample_size) if train_df is not None and self.config.sample_size and self.config.sample_size < len(train_df) else train_df
        test_sample = test_df.sample(self.config.sample_size) if test_df is not None and self.config.sample_size and self.config.sample_size < len(test_df) else test_df

        in_notebook = self._in_notebook()

        report = LeakReport(
            metadata={
                "n_rows": int(len(df_sample)),
                "n_columns": int(len(df_sample.columns)),
                "columns": list(df_sample.columns),
                "target": target,
                "time_col": time_col,
                "entity_cols": entity_cols or [],
                "has_train_df": train_sample is not None,
                "has_test_df": test_sample is not None,
                "environment": "notebook" if in_notebook else "terminal",
            }
        )

        steps = self._build_steps(
            df=df_sample,
            target=target,
            time_col=time_col,
            train_df=train_sample,
            test_df=test_sample,
            entity_cols=entity_cols,
        )

        step_results: list[dict[str, object]] = []

        progress = tqdm(
            total=len(steps),
            desc="FeatureLeak scan",
            unit="step",
            disable=not show_progress,
            dynamic_ncols=True,
            leave=True,
        )

        for check_name, description, step_fn in steps:
            if show_progress:
                progress.set_description_str(description)

            start = time.perf_counter()
            issues_before = len(report.issues)

            try:
                raw_output = step_fn()
                step_issues = self._coerce_step_output(raw_output, fallback_check=check_name)

                for issue in step_issues:
                    report.add_issue(issue)

                elapsed = time.perf_counter() - start
                issues_added = len(report.issues) - issues_before

                step_results.append(
                    {
                        "name": description,
                        "check": check_name,
                        "status": "warning" if issues_added > 0 else "ok",
                        "elapsed": elapsed,
                        "n_issues": issues_added,
                    }
                )

                if show_progress:
                    progress.set_postfix_str(f"time={elapsed:.2f}s, issues={issues_added}")

            except Exception as exc:
                elapsed = time.perf_counter() - start
                self._add_step_error(report, check_name, description, exc)

                step_results.append(
                    {
                        "name": description,
                        "check": check_name,
                        "status": "error",
                        "elapsed": elapsed,
                        "n_issues": 1,
                    }
                )

                if show_progress:
                    progress.set_postfix_str(f"time={elapsed:.2f}s, error={type(exc).__name__}")

            finally:
                progress.update(1)

        progress.close()

        self._sort_report_issues(report)
        self._finalise_risk_buckets(report)

        report.metadata["step_results"] = step_results
        report.metadata["n_issues"] = len(report.issues)
        report.metadata["n_warnings"] = len(report.issues)
        report.metadata["score"] = report.score()
        report.metadata["risk_level"] = report.risk_level()
        report.metadata["issue_sorting"] = {
            "primary": "severity_desc",
            "secondary": "check_priority_desc",
            "tertiary": "first_affected_column_desc",
            "quaternary": "title_desc",
        }

        if show_summary and not in_notebook:
            self._print_terminal_summary(report, step_results)

        return report