import gradio as gr
import joblib
from datetime import datetime
import os

from src.predict.models import (
    BaseMLModel,
    EloBaselineModel,
    LogisticRegressionModel,
    XGBoostModel,
    SVCModel,
    RandomForestModel,
    BernoulliNBModel,
    LGBMModel
)
from src.config import MODELS_DIR

# --- Model Cache ---
# This global dictionary will store loaded models to avoid reloading them from disk.
MODEL_CACHE = {}

# --- Gradio App Setup ---
if not os.path.exists(MODELS_DIR):
    os.makedirs(MODELS_DIR)
    print(f"Warning: Models directory not found. Created a dummy directory at '{MODELS_DIR}'.")

# Get a list of available models
available_models = [f for f in os.listdir(MODELS_DIR) if f.endswith(".joblib")]
if not available_models:
    print(f"Warning: No models found in '{MODELS_DIR}'. The dropdown will be empty.")
    available_models.append("No models found")

# --- Prediction Function ---
def predict_fight(model_name, fighter1_name, fighter2_name):
    """
    Loads the selected model and predicts the winner of a fight.
    """
    if model_name == "No models found" or not fighter1_name or not fighter2_name:
        return "Please select a model and enter both fighter names.", ""

    try:
        # Load model from cache or from disk if it's the first time
        if model_name not in MODEL_CACHE:
            print(f"Loading and caching model: {model_name}...")
            model_path = os.path.join(MODELS_DIR, model_name)
            MODEL_CACHE[model_name] = joblib.load(model_path)
            print("...model cached.")

        model = MODEL_CACHE[model_name]

        fight = {
            'fighter_1': fighter1_name,
            'fighter_2': fighter2_name,
            'event_date': datetime.now().strftime('%B %d, %Y')
        }

        prediction_result = model.predict(fight)

        if prediction_result and prediction_result.get('winner'):
            winner = prediction_result['winner']
            prob = prediction_result['probability']
            return winner, f"{prob:.1%}"
        else:
            return "Could not make a prediction.", ""

    except FileNotFoundError:
        return f"Error: Model file '{model_name}' not found.", ""
    except Exception as e:
        print(f"An error occurred during prediction: {e}")
        return f"An error occurred: {e}", ""

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🥋 UFC Fight Predictor 🥊")
    gr.Markdown("Select a prediction model and enter two fighter names to predict the outcome.")

    with gr.Column():
        model_dropdown = gr.Dropdown(
            label="Select Model",
            choices=available_models,
            value=available_models[0] if available_models else None
        )
        with gr.Row():
            fighter1_input = gr.Textbox(label="Fighter 1", placeholder="e.g., Jon Jones")
            fighter2_input = gr.Textbox(label="Fighter 2", placeholder="e.g., Stipe Miocic")

    predict_button = gr.Button("Predict Winner")

    with gr.Column():
        winner_output = gr.Textbox(label="Predicted Winner", interactive=False)
        prob_output = gr.Textbox(label="Confidence", interactive=False)

    predict_button.click(
        fn=predict_fight,
        inputs=[model_dropdown, fighter1_input, fighter2_input],
        outputs=[winner_output, prob_output]
    )

# --- Launch the App ---
demo.launch()