"""
src/preprocess.py

Enterprise-grade preprocessing utilities used by app.py and ModelPredictor.

Responsibilities:
- Safe date parsing and time-feature extraction
- Missing-value handling with flexible strategies
- Lightweight feature alignment to model feature list (preserving order)
- Optional basic scaling (mean/std)
- Validation helpers for incoming data
"""

from typing import List, Optional, Dict, Any
import pandas as pd
import numpy as np
import logging
import json
from pathlib import Path

logger = logging.getLogger(__name__)


class DataPrep:
    def __init__(
        self,
        fillna_strategy: str = "median",
        custom_fill: Optional[Dict[str, Any]] = None,
        apply_scaling: bool = False,
        scaling_stats: Optional[Dict[str, Dict[str, float]]] = None,
    ):
        """
        Parameters
        ----------
        fillna_strategy : 'median' or 'zero'
        custom_fill : dict {col: value} for overrides on specific columns
        apply_scaling : if True, apply standard scaling using scaling_stats
        scaling_stats : dict {col: {"mean": .., "std": ..}}
        """
        if fillna_strategy not in ("median", "zero"):
            raise ValueError("fillna_strategy must be 'median' or 'zero'")

        self.fillna_strategy = fillna_strategy
        self.custom_fill = custom_fill or {}
        self.apply_scaling = apply_scaling
        self.scaling_stats = scaling_stats or {}

    # ---------------------------------------------------------
    # INTERNAL HELPERS
    # ---------------------------------------------------------

    def _parse_dates(self, df: pd.DataFrame) -> pd.DataFrame:
        """If a date column exists, generate useful time-based features."""
        df = df.copy()

        date_candidates = [
            c for c in df.columns
            if c.lower() in ("date", "timestamp", "datetime")
        ]
        if not date_candidates:
            return df

        date_col = date_candidates[0]

        try:
            df[date_col] = pd.to_datetime(df[date_col], errors="coerce")
            df["dayofweek"] = df[date_col].dt.dayofweek
            df["month"] = df[date_col].dt.month

            # Only add hour if it varies
            if "hour" not in df.columns:
                if df[date_col].dt.hour.nunique(dropna=True) > 1:
                    df["hour"] = df[date_col].dt.hour
        except Exception as e:
            logger.debug(f"Date parsing failed for column {date_col}: {e}")

        return df

    def _fill_missing(self, df: pd.DataFrame) -> pd.DataFrame:
        """Fill missing values with strategy + custom overrides."""
        df = df.copy()

        # apply custom fills first
        for col, val in self.custom_fill.items():
            if col in df.columns:
                try:
                    df[col] = df[col].fillna(val)
                except Exception:
                    df[col] = df[col].astype(object).fillna(val)

        # numeric fill
        numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()

        if numeric_cols:
            if self.fillna_strategy == "median":
                medians = df[numeric_cols].median()
                df[numeric_cols] = df[numeric_cols].fillna(medians)
            else:  # zero fill
                df[numeric_cols] = df[numeric_cols].fillna(0)

        # object/categorical fill
        obj_cols = df.select_dtypes(include=["object", "category"]).columns.tolist()
        if obj_cols:
            df[obj_cols] = df[obj_cols].fillna("")

        return df

    def _apply_scaling(self, X: pd.DataFrame) -> pd.DataFrame:
        """Standard scaling (X - mean) / std."""
        if not self.apply_scaling or not self.scaling_stats:
            return X

        X = X.copy()

        for col, stats in self.scaling_stats.items():
            if col in X.columns:
                mean = stats.get("mean", 0.0)
                std = stats.get("std", 1.0) or 1.0
                try:
                    X[col] = (X[col] - mean) / std
                except Exception as e:
                    logger.debug(f"Scaling failed for {col}: {e}")

        return X

    def _align_columns(self, df: pd.DataFrame, feature_columns: List[str]) -> pd.DataFrame:
        """
        Ensures the output has exactly the feature_columns, in order.
        Missing columns are created as zero; extra columns are dropped.
        """
        aligned = pd.DataFrame(index=df.index)

        for col in feature_columns:
            if col in df.columns:
                aligned[col] = df[col]
            else:
                aligned[col] = 0

        return aligned

    # ---------------------------------------------------------
    # PUBLIC API
    # ---------------------------------------------------------

    def validate(self, df: pd.DataFrame) -> None:
        """Lightweight data checks."""
        if df is None or len(df) == 0:
            raise ValueError("Input DataFrame is empty.")
        if not isinstance(df, pd.DataFrame):
            raise ValueError("Input must be a pandas DataFrame.")

    def prepare(self, df: pd.DataFrame, feature_columns: List[str]) -> pd.DataFrame:
        """
        Main preprocessing entrypoint.

        Steps:
        1. Validate
        2. Date parsing
        3. Missing-value handling
        4. Optional scaling
        5. Alignment with feature list
        """
        self.validate(df)

        df = df.copy()
        df = self._parse_dates(df)
        df = self._fill_missing(df)
        X = self._align_columns(df, feature_columns)

        if self.apply_scaling:
            X = self._apply_scaling(X)

        return X[feature_columns]  # ensure strict order

    def save_scaling_stats(self, out_path: str, feature_cols: List[str], df_reference: pd.DataFrame) -> None:
        """
        Computes & saves scaling means/stds for future runs.
        """
        stats = {}
        numeric = df_reference[feature_cols].select_dtypes(include=[np.number]).columns.tolist()

        for col in numeric:
            stats[col] = {
                "mean": float(df_reference[col].mean()),
                "std": float(df_reference[col].std() or 1.0)
            }

        Path(out_path).write_text(json.dumps(stats, indent=2))