"""
Model Trainer Module - Fixed
- Proper train/val/holdout split
- early_stopping_rounds actually passed to model.fit()
- Native XGBoost .ubj format (not pickle)
- No StandardScaler for tree models
- 5-fold cross-validation
- No Streamlit
"""
import os
import joblib
import xgboost as xgb
import pandas as pd
from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
from sklearn.metrics import roc_auc_score, f1_score, precision_score, recall_score, confusion_matrix, roc_curve
from imblearn.over_sampling import SMOTE

from src.config_loader import get_config


def _get_feature_cols() -> list:
    return get_config()['ml']['feature_cols']


def train_model(feature_df: pd.DataFrame) -> tuple:
    cfg = get_config()['ml']
    FEATURE_COLS = _get_feature_cols()

    # Filter to only columns that exist in the DataFrame
    available = [c for c in FEATURE_COLS if c in feature_df.columns]
    X = feature_df[available].copy()
    y = feature_df['fraud_flag'].copy()

    # FIX: proper 70/15/15 split
    X_temp, X_hold, y_temp, y_hold = train_test_split(
        X, y, test_size=cfg['val_size'], random_state=cfg['random_state'], stratify=y
    )
    X_train, X_val, y_train, y_val = train_test_split(
        X_temp, y_temp, test_size=cfg['val_size'], random_state=cfg['random_state'], stratify=y_temp
    )

    pos = (y_train == 1).sum()
    neg = (y_train == 0).sum()
    scale_pos_weight = neg / max(pos, 1)

    # Optional SMOTE for severe class imbalance
    if pos / max(len(y_train), 1) < 0.05:
        try:
            sm = SMOTE(random_state=cfg['random_state'])
            X_train, y_train = sm.fit_resample(X_train, y_train)
        except Exception:
            pass

    model = xgb.XGBClassifier(
        n_estimators=cfg['n_estimators'],
        max_depth=cfg['max_depth'],
        learning_rate=cfg['learning_rate'],
        subsample=cfg['subsample'],
        colsample_bytree=cfg['colsample'],
        scale_pos_weight=scale_pos_weight,
        eval_metric='auc',
        random_state=cfg['random_state'],
        tree_method='hist',
        early_stopping_rounds=cfg['early_stopping'],  # FIX: actually passed now
    )

    model.fit(
        X_train, y_train,
        eval_set=[(X_val, y_val)],
        verbose=False,
    )

    y_pred = model.predict(X_hold)
    y_proba = model.predict_proba(X_hold)[:, 1]
    fpr, tpr, _ = roc_curve(y_hold, y_proba)

    # 5-fold CV on full dataset
    cv_model = xgb.XGBClassifier(
        n_estimators=model.best_iteration + 1 if hasattr(model, 'best_iteration') else 100,
        max_depth=cfg['max_depth'],
        learning_rate=cfg['learning_rate'],
        tree_method='hist',
    )
    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=cfg['random_state'])
    cv_scores = cross_val_score(cv_model, X, y, cv=skf, scoring='roc_auc')

    metrics = {
        'auc_roc': float(roc_auc_score(y_hold, y_proba)),
        'f1': float(f1_score(y_hold, y_pred, zero_division=0)),
        'precision': float(precision_score(y_hold, y_pred, zero_division=0)),
        'recall': float(recall_score(y_hold, y_pred, zero_division=0)),
        'confusion_matrix': confusion_matrix(y_hold, y_pred).tolist(),
        'fpr': fpr.tolist(),
        'tpr': tpr.tolist(),
        'cv_auc_mean': float(cv_scores.mean()),
        'cv_auc_std': float(cv_scores.std()),
        'feature_cols': available,
        'n_features': len(available),
    }

    # FIX: save in native XGBoost format + joblib for feature list
    os.makedirs(os.path.dirname(cfg['model_path']), exist_ok=True)
    model.save_model(cfg['model_path'])
    joblib.dump({'feature_cols': available}, cfg['scaler_path'])

    return model, metrics