import os
import pandas as pd
import joblib
from sklearn.ensemble import RandomForestRegressor
from xgboost import XGBRegressor
from sklearn.metrics import mean_absolute_error
from config.database import get_db_connection
from services.feature_builder import prepare_features, FEATURE_COLUMNS

MODEL_DIR = "models"
os.makedirs(MODEL_DIR, exist_ok=True)

def train_and_save_models():
    conn = get_db_connection()
    
    # Query gabungan antara tabel transaksi dan transaksi_detail
    query = """
        SELECT 
            td.produk_id,
            td.jumlah AS target_jumlah,
            t.tanggal,
            t.bulan,
            t.tahun,
            t.is_weekend,
            t.is_libur_nasional,
            t.is_ramadhan,
            t.is_idul_fitri,
            t.is_idul_adha,
            t.is_tahun_baru
        FROM transaksi_detail td
        JOIN transaksi t ON td.transaksi_id = t.id
    """
    
    df = pd.read_sql(query, conn)
    conn.close()
    
    if df.empty:
        return {"status": "error", "message": "Data transaksi masih kosong!"}
        
    # Preprocessing
    df = prepare_features(df)
    
    X = df[FEATURE_COLUMNS]
    y = df['target_jumlah']
    
    # Split manual/sederhana atau langsung train seluruh histori
    # 1. Melatih Model Random Forest
    rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
    rf_model.fit(X, y)
    rf_pred = rf_model.predict(X)
    mae_rf = mean_absolute_error(y, rf_pred)
    
    # 2. Melatih Model XGBoost
    xgb_model = XGBRegressor(n_estimators=100, learning_rate=0.1, random_state=42)
    xgb_model.fit(X, y)
    xgb_pred = xgb_model.predict(X)
    mae_xgb = mean_absolute_error(y, xgb_pred)
    
    # 3. Simpan Model Terlatih ke File .pkl
    joblib.dump(rf_model, os.path.join(MODEL_DIR, "model_rf.pkl"))
    joblib.dump(xgb_model, os.path.join(MODEL_DIR, "model_xgb.pkl"))
    
    return {
        "status": "success",
        "message": "Model berhasil dilatih dan disimpan!",
        "mae_random_forest": float(mae_rf),
        "mae_xgboost": float(mae_xgb)
    }