import gradio as gr
import pickle
import pandas as pd
import traceback
import json
import xgboost as xgb

# 1. Load bundle model
with open("best_model_xgb.pkl", "rb") as f:
    bundle = pickle.load(f)

model = bundle['model']
impute_col = bundle['impute_col']
feature_names = bundle['feature_names']

# 2. Fungsi Prediksi dengan Print ke Log Hugging Face
def predict_kidney(input_data):
    try:
        # --- MULAI SESI DEBUGGING KE LOG HUGGING FACE ---
        print("\n" + "="*40, flush=True)
        print("[INFO] Menerima Request API Baru", flush=True)
        
        if isinstance(input_data, str):
            input_data = json.loads(input_data)
            
        received_keys = list(input_data.keys()) if isinstance(input_data, dict) else []
        missing_keys = [col for col in feature_names if col not in received_keys]
        
        # Cetak detail ke terminal/log container
        print(f"[DEBUG] Total Parameter Diharapkan: {len(feature_names)}", flush=True)
        print(f"[DEBUG] Total Parameter Diterima  : {len(received_keys)}", flush=True)
        
        if missing_keys:
            print(f"[WARNING] Parameter Kurang: {missing_keys}", flush=True)
        else:
            print("[SUCCESS] Semua parameter lengkap!", flush=True)
            
        print("="*40 + "\n", flush=True)
        # --- AKHIR SESI DEBUGGING ---

        # Proses Data
        df_input = pd.DataFrame([input_data])
        
        for col in feature_names:
            if col not in df_input.columns:
                df_input[col] = 0.0
                
        for col in impute_col:
            if pd.isna(df_input[col].iloc[0]):
                df_input[col] = 0.0
                
        df_input = df_input[feature_names]
        
        # Prediksi
        prediction = model.predict(df_input)[0]
        
        return {
            "status": "success",
            "prediction_class": int(prediction)
        }
    except Exception as e:
        # Print error ke log Hugging Face juga jika terjadi crash
        print(f"\n[ERROR] {str(e)}", flush=True)
        print(traceback.format_exc(), flush=True)
        
        return {
            "status": "error",
            "message": str(e),
            "traceback": traceback.format_exc()
        }

# 3. Setup UI Gradio
demo = gr.Interface(
    fn=predict_kidney,
    inputs=gr.JSON(label="Input Features (JSON)"),
    outputs=gr.JSON(label="Prediction Output"),
    title="NutriSnapS Kidney Prevention API (XGBoost)",
    api_name="predict_kidney"
)

if __name__ == "__main__":
    demo.launch(show_error=True, ssr_mode=False)