from typing import Optional, Dict
import pandas as pd
import numpy as np
from quickml.artifact import ArtifactManager

from sklearn.model_selection import train_test_split

from quickml.decisions import Decision, DecisionLog
from quickml.schema import SchemaInferer
from quickml.preprocess import DataQualityChecker, Preprocessor
from quickml.modeling import ModelTrainer


class Pipeline:
    def __init__(
        self,
        model: str,
        task: Optional[str] = None,
        test_size: float = 0.2,
        random_state: int = 42,
        strict: bool = False,
    ):
        self.model_name = model
        self.task = task
        self.test_size = test_size
        self.random_state = random_state
        self.strict = strict

        self.log = DecisionLog()

        self.schema = None
        self.preprocessor = None
        self.trainer = None
        self.metrics = None
        self.target = None

    def fit(self, data: pd.DataFrame, target: str) -> Dict[str, float]:
        self.target = target

        X_df = data.drop(columns=[target])
        y = data[target].values

        # Infer task if not provided
        if self.task is None:
            self.task = (
                "regression"
                if pd.api.types.is_numeric_dtype(data[target])
                else "classification"
            )

        # 1. Schema inference
        inferer = SchemaInferer(self.log, strict=self.strict)
        self.schema = inferer.infer(X_df)

        # 2. Data quality checks
        checker = DataQualityChecker(
            log=self.log,
            target=target,
            task=self.task,
            strict=self.strict,
        )
        checker.run(data)

        # 3. Preprocessing
        self.preprocessor = Preprocessor(self.schema, self.log)
        self.preprocessor.fit(X_df)
        X = self.preprocessor.transform(X_df)

        # 4. Train/test split
        X_train, X_test, y_train, y_test = train_test_split(
            X,
            y,
            test_size=self.test_size,
            random_state=self.random_state,
        )

        # 5. Model training
        self.trainer = ModelTrainer(
            model_name=self.model_name,
            task=self.task,
            log=self.log,
            random_state=self.random_state,
        )
        self.trainer.fit(X_train, y_train)

        # 6. Evaluation
        self.metrics = self.trainer.evaluate(X_test, y_test)
        return self.metrics

    def predict(self, data: pd.DataFrame):
        if self.preprocessor is None or self.trainer is None:
            raise RuntimeError("Pipeline has not been fitted or loaded.")

        trained_cols = set(self.schema.keys())
        input_cols = set(data.columns)

        missing = trained_cols - input_cols
        extra = input_cols - trained_cols

        if missing or extra:
            self.log.add(
                Decision.create(
                    category="inference",
                    component="Pipeline",
                    action="schema_mismatch",
                    reason="Input schema does not match training schema",
                    parameters={
                        "missing_columns": list(missing),
                        "extra_columns": list(extra),
                    },
                    impact="Prediction may be invalid",
                    severity="error",
                )
            )
            raise ValueError("Schema mismatch during prediction.")

        X = self.preprocessor.transform(data)
        return self.trainer.model.predict(X)

    def report(self) -> Dict[str, list]:
        return {
            "info": self.log.by_severity("info"),
            "warnings": self.log.by_severity("warning"),
            "errors": self.log.by_severity("error"),
        }
    def save(self, path: str) -> None:
        if self.trainer is None or self.preprocessor is None:
            raise RuntimeError("Pipeline must be fitted before saving.")

        payload = {
            "version": ArtifactManager.VERSION,
            "model_name": self.model_name,
            "task": self.task,
            "schema": self.schema,
            "preprocessor": self.preprocessor,
            "model": self.trainer.model,
            "decision_log": self.log,
        }

        ArtifactManager.save(path, payload)

    @classmethod
    def load(cls, path: str) -> "Pipeline":
        payload = ArtifactManager.load(path)

        pipe = cls(
            model=payload["model_name"],
            task=payload["task"],
        )

        pipe.schema = payload["schema"]
        pipe.preprocessor = payload["preprocessor"]
        pipe.log = payload["decision_log"]

        # Reattach trained model
        pipe.trainer = type("LoadedTrainer", (), {})()
        pipe.trainer.model = payload["model"]

        return pipe

