from typing import Dict
import pandas as pd

from quickml.decisions import Decision, DecisionLog


class DataQualityChecker:
    def __init__(
        self,
        log: DecisionLog,
        target: str,
        task: str,
        strict: bool = False,
    ):
        self.log = log
        self.target = target
        self.task = task
        self.strict = strict

    def run(self, df: pd.DataFrame) -> None:
        self._check_dataset_size(df)
        self._check_constant_columns(df)
        self._check_missing_values(df)

        if self.task == "classification":
            self._check_class_imbalance(df)

        if self.strict and self.log.has_errors():
            raise RuntimeError("Critical data quality errors detected.")

    def _check_dataset_size(self, df: pd.DataFrame) -> None:
        n_rows = len(df)

        if n_rows < 30:
            severity = "error"
        elif n_rows < 100:
            severity = "warning"
        else:
            return

        self.log.add(
            Decision.create(
                category="data_quality",
                component="DataQualityChecker",
                action="check_dataset_size",
                reason=f"Dataset contains only {n_rows} rows",
                parameters={"rows": n_rows},
                impact="Model may overfit or metrics may be unreliable",
                severity=severity,
            )
        )

    def _check_constant_columns(self, df: pd.DataFrame) -> None:
        for column in df.columns:
            unique_vals = df[column].dropna().unique()

            if len(unique_vals) <= 1:
                self.log.add(
                    Decision.create(
                        category="data_quality",
                        component="DataQualityChecker",
                        action="constant_column",
                        reason="Column has a single unique value",
                        parameters={"column": column},
                        impact="Provides no predictive value",
                        severity="warning",
                    )
                )

    def _check_missing_values(self, df: pd.DataFrame) -> None:
        for column in df.columns:
            missing_ratio = df[column].isna().mean()

            if missing_ratio > 0.9:
                severity = "error"
            elif missing_ratio > 0.5:
                severity = "warning"
            else:
                continue

            self.log.add(
                Decision.create(
                    category="data_quality",
                    component="DataQualityChecker",
                    action="high_missing_ratio",
                    reason=f"{missing_ratio:.0%} values missing",
                    parameters={"column": column, "missing_ratio": missing_ratio},
                    impact="May distort learning or require aggressive imputation",
                    severity=severity,
                )
            )

    def _check_class_imbalance(self, df: pd.DataFrame) -> None:
        counts = df[self.target].value_counts(normalize=True)

        if counts.iloc[0] > 0.9:
            self.log.add(
                Decision.create(
                    category="data_quality",
                    component="DataQualityChecker",
                    action="class_imbalance",
                    reason="One class dominates the target distribution",
                    parameters={"distribution": counts.to_dict()},
                    impact="Model may be biased toward majority class",
                    severity="warning",
                )
            )
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline as SkPipeline


class Preprocessor:
    def __init__(
        self,
        schema: Dict[str, str],
        log: DecisionLog,
    ):
        self.schema = schema
        self.log = log
        self.transformer = None

    def fit(self, df: pd.DataFrame) -> None:
        numeric_cols = [c for c, t in self.schema.items() if t == "numeric"]
        categorical_cols = [c for c, t in self.schema.items() if t == "categorical"]

        # Numeric pipeline
        numeric_pipeline = SkPipeline(
            steps=[
                ("imputer", SimpleImputer(strategy="median")),
            ]
        )

        for col in numeric_cols:
            missing_ratio = df[col].isna().mean()
            self.log.add(
                Decision.create(
                    category="preprocessing",
                    component="Preprocessor",
                    action="numeric_imputation",
                    reason="Median imputation applied",
                    parameters={
                        "column": col,
                        "strategy": "median",
                        "missing_ratio": missing_ratio,
                    },
                    impact="Missing numeric values filled",
                    severity="info",
                )
            )

        # Categorical pipeline
        categorical_pipeline = SkPipeline(
            steps=[
                ("imputer", SimpleImputer(strategy="constant", fill_value="__missing__")),
                (
                    "encoder",
                    OneHotEncoder(
                        handle_unknown="ignore",
                        sparse_output=False,
                    ),
                ),
            ]
        )

        for col in categorical_cols:
            self.log.add(
                Decision.create(
                    category="preprocessing",
                    component="Preprocessor",
                    action="categorical_encoding",
                    reason="One-hot encoding applied",
                    parameters={"column": col, "encoding": "onehot"},
                    impact="Categorical expanded into binary features",
                    severity="info",
                )
            )

        self.transformer = ColumnTransformer(
            transformers=[
                ("num", numeric_pipeline, numeric_cols),
                ("cat", categorical_pipeline, categorical_cols),
            ],
            remainder="drop",
        )

        self.transformer.fit(df)

    def transform(self, df: pd.DataFrame):
        if self.transformer is None:
            raise RuntimeError("Preprocessor has not been fitted.")

        return self.transformer.transform(df)
