import gradio as gr
import torch
from pathlib import Path
import traceback
import sys

# Force reload of model module
if 'model' in sys.modules:
    del sys.modules['model']

from model import miniGPT

DEVICE = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
print(f"Loading model on device: {DEVICE}")

SCRIPT_DIR = Path(__file__).parent
model_dir = SCRIPT_DIR / "trained_model_small"

trainer = None
error_message = None

if model_dir.exists():
    print(f"Loading model from {model_dir}")
    print(f"Model directory contents:")
    for item in model_dir.iterdir():
        print(f"  - {item.name}")
    try:
        trainer = miniGPT.load(str(model_dir), map_location=DEVICE)
        trainer = trainer.to(DEVICE)
        print("✓ Model loaded successfully!")
        print(f"✓ Model ready for inference on {DEVICE}!")
    except Exception as e:
        error_message = str(e)
        print(f"✗ Error loading model: {e}")
        print(f"✗ Full traceback:")
        print(traceback.format_exc())
        trainer = None
else:
    error_message = f"Model directory not found at {model_dir}"
    print(f"✗ {error_message}")
    print(f"✗ Contents of {SCRIPT_DIR}:")
    for item in SCRIPT_DIR.iterdir():
        print(f"  - {item.name}")


def generate_recipe(prompt, max_length=300, temperature=0.85, top_k=50, top_p=0.95):
    if trainer is None:
        if error_message:
            return f"❌ Error: Model not loaded.\n\nDetails: {error_message}"
        return "❌ Error: Model not loaded"
    
    try:
        generated = trainer.generate(
            prompt,
            max_length=max_length,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
        )
        return generated
    except Exception as e:
        return f"❌ Error generating recipe: {str(e)}"

with gr.Blocks(title="Recipe Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Welcome eggs 1/4 cup - stupidcookgpt
        
        Generate recipes using a fine-tuned miniGPT model trained on recipe data.
        
        Verified to generate your next delicious meal in 20 seconds (unless it doesn't).
        
        **Model Info:** ~124M parameters, trained on corbt/all-recipes.

        *small print*: ppl this is a friggin demo not some gordon ramsay, pls use common sense and don't sue me if you eat its egg-chocolate-broccoli surprise.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Recipe Prompt",
                placeholder="e.g., Recipe for chocolate cake:\n\nIngredients:",
                lines=3,
                value="Recipe for chocolate ",
            )
        with gr.Column():
            generate_btn = gr.Button("Tell the chef...", variant="primary", size="lg")
    
    output = gr.Textbox(label="Generated Recipe", lines=12, max_lines=20)
    
    with gr.Row():
        with gr.Column():
            max_length = gr.Slider(50, 500, 300, step=50, label="Max Length", info="Maximum tokens to generate")
        with gr.Column():
            temperature = gr.Slider(0.1, 2.0, 0.85, step=0.1, label="Temperature", info="Higher = more creative")
    
    with gr.Row():
        with gr.Column():
            top_k = gr.Slider(0, 100, 50, step=5, label="Top-K", info="Sample from top K tokens")
        with gr.Column():
            top_p = gr.Slider(0.0, 1.0, 0.95, step=0.05, label="Top-P", info="Nucleus sampling threshold")
    
    generate_btn.click(
        fn=generate_recipe,
        inputs=[prompt_input, max_length, temperature, top_k, top_p],
        outputs=output,
    )
    
    # example prompts
    gr.Examples(
        examples=[
            ["Recipe for chocolate ", 300, 0.85, 50, 0.95],
            ["Recipe for a delicious apple pie:\n\nIngredients:\n", 300, 0.85, 50, 0.95],
            ["Ingredients: flour butter sugar eggs", 300, 0.85, 50, 0.95],
            ["Instructions: preheat oven to 350", 300, 0.85, 50, 0.95],
            ["How to make a perfect coffee cake", 300, 0.85, 50, 0.95],
        ],
        inputs=[prompt_input, max_length, temperature, top_k, top_p],
        outputs=output,
        fn=generate_recipe,
        cache_examples=False,
    )
    
    gr.Markdown(
        """
        ### How to use:
        1. Enter a prompt 4 the chef
        2. Adjust chef parameters if desired (make it gordon ramsay when Temp=2 or a chill cook with Temp=0.4)
        3. Click "Tell the chef..." to create a new recipe
        
        ### Tips:
        - **Temperature**: Higher values (1.0+) make the chef more creative but probably more sus.
        - **Top-K**: Limits sampling to top K most likely tokens
        - **Top-P**: Only considers tokens with cumulative probability <= P
        """
    )


if __name__ == "__main__":
    demo.launch()