import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os # For Hugging Face token

# Import spaces for ZeroGPU if you need to decorate specific functions
# For models loaded via transformers and run on a device managed by ZeroGPU,
# explicit @spaces.GPU might not always be needed directly on the inference function
# if the entire space is on ZeroGPU hardware. However, for clarity or complex setups:
# import spaces # Uncomment if using @spaces.GPU decorator

# --- Configuration ---
HF_TOKEN = os.getenv("HF_TOKEN") # Recommended to store your Hugging Face token as a Space secret

MODEL_OPTIONS = {
    "Qwen1.5-1.8B-Chat": "Qwen/Qwen1.5-1.8B-Chat",
    "Qwen2.5-Coder-3B": "Qwen/Qwen2.5-Coder-3B", # Example for a Qwen code model around 3B params
}

# --- Model Loading Cache ---
# This dictionary will cache loaded models and tokenizers to avoid reloading on every call
loaded_models = {}

def get_model_and_tokenizer(model_name_key):
    if model_name_key not in loaded_models:
        model_id = MODEL_OPTIONS[model_name_key]
        print(f"Loading model: {model_id}...")
        try:
            # Ensure you have accepted the terms of use for these models on Hugging Face Hub
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                torch_dtype="auto", # Let transformers decide the best dtype
                device_map="auto",  # Automatically maps model to available device (GPU on ZeroGPU)
                token=HF_TOKEN # Use token if model is private or requires it
            )
            tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
            loaded_models[model_name_key] = (model, tokenizer)
            print(f"Model {model_id} loaded successfully.")
        except Exception as e:
            print(f"Error loading model {model_id}: {e}")
            # Fallback or error handling
            if model_name_key in loaded_models: # Remove if partially loaded
                del loaded_models[model_name_key]
            raise gr.Error(f"Failed to load model {model_name_key}. Please check the model ID and your Hugging Face token permissions. Error: {e}")
    return loaded_models[model_name_key]

# --- Inference Function ---
# If you need finer-grained control over GPU allocation for specific parts:
# @spaces.GPU(duration=120) # Example: Request GPU for 120 seconds for this function
def generate_response(prompt_text, model_choice, max_new_tokens=512, temperature=0.7, top_p=0.9):
    if not prompt_text:
        return "Please enter a prompt."
    if not model_choice:
        return "Please select a model."

    try:
        model, tokenizer = get_model_and_tokenizer(model_choice)
    except Exception as e:
        return str(e) # Display loading error to user

    device = model.device # Get the device the model is on

    if "Chat" in model_choice: # Apply chat template for chat models
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt_text}
        ]
        try:
            input_text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
        except Exception as e: # Fallback if apply_chat_template has issues or is not applicable
            print(f"Warning: Could not apply chat template for {model_choice}: {e}. Using prompt as is.")
            input_text = prompt_text

    else: # For code or non-chat models, use the prompt directly or adjust as needed
        input_text = prompt_text

    model_inputs = tokenizer([input_text], return_tensors="pt").to(device)

    try:
        generated_ids = model.generate(
            model_inputs.input_ids,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True # Necessary for temperature and top_p to have an effect
        )
        # For some models, the input prompt is included in the generated_ids.
        # We need to decode only the newly generated tokens.
        # This slicing can vary based on the model and tokenizer.
        # A common approach is to slice based on the input_ids length:
        response_ids = generated_ids[0][model_inputs.input_ids.shape[-1]:]
        response_text = tokenizer.decode(response_ids, skip_special_tokens=True)

    except Exception as e:
        print(f"Error during generation with {model_choice}: {e}")
        return f"Error generating response: {e}"

    return response_text

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# LLM Coding & Math Experiment")
    gr.Markdown("Query Qwen1.5-1.8B-Chat or Qwen Code models using ZeroGPU.")

    with gr.Row():
        model_dropdown = gr.Dropdown(
            label="Select Model",
            choices=list(MODEL_OPTIONS.keys()),
            value=list(MODEL_OPTIONS.keys())[0] # Default to the first model
        )
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter your prompt:", lines=4, placeholder="e.g., Write a Python function to calculate factorial, or What is the capital of France?")
    with gr.Row():
        output_text = gr.Textbox(label="Model Response:", lines=8, interactive=False)

    with gr.Row():
        submit_button = gr.Button("Generate Response", variant="primary")

    with gr.Accordion("Advanced Settings", open=False):
        max_new_tokens_slider = gr.Slider(minimum=32, maximum=2048, value=512, step=32, label="Max New Tokens")
        temperature_slider = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.05, label="Temperature")
        top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P")


    # Event listener for the button
    submit_button.click(
        fn=generate_response,
        inputs=[prompt_input, model_dropdown, max_new_tokens_slider, temperature_slider, top_p_slider],
        outputs=output_text,
        api_name="generate" # Exposes an API endpoint
    )

    gr.Markdown("## Notes:")
    gr.Markdown(
        "- Ensure you have accepted the terms of use for the selected Qwen models on the Hugging Face Hub.\n"
        "- Model loading can take some time, especially on the first run or when switching models.\n"
        "- This Space runs on ZeroGPU, which means GPU resources are allocated dynamically."
    )

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