"""preprocess_data handler — scale, encode, impute, drop columns."""
from typing import Optional

import polars as pl

from ..session_store import SessionStore, _generate_id
from ._helpers import _resolve_dataset


async def _handle_preprocess_data(
    dataset_id: str,
    columns: list[str],
    config: Optional[dict] = None,
) -> dict:
    """Preprocess selected columns (scale, encode, impute, drop)."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err

    lf: pl.LazyFrame = entry['lazyframe']
    cfg = config or {}
    column_mapping: dict[str, str] = {}

    df = lf.collect(streaming=True)

    # Resolve columns that actually exist
    available = set(df.columns)
    target_cols = [c for c in columns if c in available]
    if not target_cols:
        return {'error': 'None of the specified columns exist in the dataset',
                'truncated': False}

    # Drop columns
    drop_cols = cfg.get('drop', [])
    if drop_cols:
        df = df.drop([c for c in drop_cols if c in df.columns])
        for c in drop_cols:
            if c in available:
                column_mapping[c] = f'<dropped:{c}>'

    # Fill missing values
    fill_strategy = cfg.get('fillna')
    if fill_strategy and fill_strategy != 'drop':
        for col in target_cols:
            if col not in df.columns:
                continue
            if df[col].null_count() == 0:
                continue
            if fill_strategy == 'mean':
                mean_val = df[col].drop_nulls().mean()
                if mean_val is not None:
                    df = df.with_columns(pl.col(col).fill_null(mean_val))
            elif fill_strategy == 'median':
                med_val = df[col].drop_nulls().median()
                if med_val is not None:
                    df = df.with_columns(pl.col(col).fill_null(med_val))
            elif fill_strategy == 'zero':
                df = df.with_columns(pl.col(col).fill_null(0))
    elif fill_strategy == 'drop':
        df = df.drop_nulls(subset=target_cols)

    # Standard scaling (z-score)
    if cfg.get('standardize'):
        from sklearn.preprocessing import StandardScaler
        num_cols = [c for c in target_cols if c in df.columns
                    and df[c].dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64)]
        if num_cols:
            scaler = StandardScaler()
            scaled = scaler.fit_transform(df.select(num_cols).to_numpy())
            for i, col in enumerate(num_cols):
                col_scaled = f'{col}_scaled'
                df = df.with_columns(pl.Series(col_scaled, scaled[:, i]))
                column_mapping[col] = col_scaled

    # One-hot encoding
    if cfg.get('onehot'):
        str_cols = [c for c in target_cols if c in df.columns
                    and df[c].dtype in (pl.Utf8, pl.String, pl.Categorical)]
        if str_cols:
            df = df.to_dummies(columns=str_cols, drop_first=False)
            for c in str_cols:
                column_mapping[c] = f'<onehot:{c}>'

    new_id = _generate_id('pp')
    new_lf = df.lazy()
    new_schema = {name: str(dtype) for name, dtype in
                  zip(new_lf.collect_schema().names(), new_lf.collect_schema().dtypes())}

    store = SessionStore()
    store.store_dataset(
        dataset_id=new_id,
        lazyframe=new_lf,
        schema=new_schema,
        metadata={
            'row_count': len(df),
            'parent_dataset_id': dataset_id,
            'preprocessing_config': cfg,
        },
        parent_id=dataset_id,
    )

    return {
        'preprocessed_dataset_id': new_id,
        'column_mapping': column_mapping,
        'row_count': len(df),
        'column_count': len(new_schema),
    }
