# ==============================================================================
# V5 GRADIO DEPLOYMENT SCRIPT (UPGRADED UI)
# ==============================================================================
# This script creates a beautiful, user-friendly web UI for your v5 model.
#
# TO DEPLOY ON HUGGING FACE SPACES:
# 1. Create a new Space and choose the "Gradio" SDK.
# 2. Select the free "CPU basic" hardware.
# 3. Create a file named `app.py` and paste this code into it.
# 4. Create a `requirements.txt` file and add the libraries listed below.
# ==============================================================================

# requirements.txt file contents:
# gradio
# transformers
# peft
# accelerate
# bitsandbytes
# torch

import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# --- Configuration ---
# Your final, v5 model on the Hugging Face Hub
MODEL_ID = "Arko007/my-awesome-code-assistant-v5" 
BASE_MODEL_ID = "codellama/CodeLlama-7b-hf"

# --- Load the Model (Memory-Optimized) ---
print("Setting up 4-bit quantization...")
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)

print(f"Loading fine-tuned model: {MODEL_ID}...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token

print("✅ Model loaded successfully!")

# --- Inference Function ---
def generate_code(instruction, progress=gr.Progress(track_tqdm=True)):
    """
    Generates code from an instruction using the v5 model.
    """
    progress(0, desc="Formatting prompt...")
    prompt = f"""### Instruction:
{instruction}

### Response:"""
    
    # *** FIX: Use model.device to automatically select CPU or GPU ***
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    progress(0.2, desc="Generating tokens...")
    outputs = model.generate(
        **inputs,
        max_new_tokens=2048,
        temperature=0.1,
        top_p=0.9,
        eos_token_id=tokenizer.eos_token_id
    )
    
    progress(0.8, desc="Decoding response...")
    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    code_part = response_text.split("### Response:")[1].strip()
    return code_part

# --- Create and Launch the Gradio Web App ---
print("Launching Gradio app with upgraded UI...")

# *** FIX: Add custom CSS and a better theme ***
css = """
body { background-color: #0F172A; }
.gradio-container { max-width: 800px !important; margin: auto !important; }
footer { display: none !important; }
"""

# *** FIX: Use a modern theme ***
theme = gr.themes.Glass(
    primary_hue="sky", 
    secondary_hue="blue",
    neutral_hue="slate"
).set(
    body_background_fill="#0F172A",
    block_background_fill="#1E293B",
    block_border_width="1px",
    block_title_background_fill="none",
    input_background_fill="#0F172A",
)

with gr.Blocks(theme=theme, css=css) as demo:
    gr.Markdown("# 🤖 My Awesome Code Assistant (v5)")
    gr.Markdown("### Built by Arko007. Powered by a custom fine-tuned Code Llama model.")
    
    with gr.Column():
        instruction_box = gr.Textbox(
            lines=5, 
            label="Instruction", 
            placeholder="e.g., Write a Python script to scrape a website for all its links.",
            elem_id="instruction-textbox" # Add an ID for our JavaScript
        )
        output_box = gr.Code(label="Generated Code", language="python", interactive=False)
        
    generate_button = gr.Button("Generate Code", variant="primary", elem_id="generate-button")
    
    gr.Markdown("---")
    gr.Markdown("### Or, try one of these examples:")
    examples = gr.Examples(
        examples=[
            "Write a Python function to find the factorial of a number using recursion.",
            "Create a simple Flask API with a single endpoint that returns 'Hello, World!'.",
            "Write a C++ program to implement a binary search tree.",
            "Explain the concept of closures in JavaScript with a code example."
        ],
        inputs=instruction_box
    )
    
    # Connect the button click to the function
    generate_button.click(
        fn=generate_code, 
        inputs=instruction_box, 
        outputs=output_box
    )

    # *** FIX: Add JavaScript to make "Enter" submit the form ***
    # This JS listens for a keypress on our textbox. If the key is Enter
    # and the Shift key is NOT held down, it clicks the generate button.
    js_code = """
    <script>
    function onKeyPress(event) {
        var instructionTextbox = document.getElementById('instruction-textbox').querySelector('textarea');
        if (event.key === 'Enter' && !event.shiftKey) {
            event.preventDefault(); // Prevent new line
            var generateButton = document.getElementById('generate-button');
            generateButton.click(); // Click the button
        }
    }
    // We need to wait for the Gradio app to mount the elements
    document.addEventListener('DOMContentLoaded', function() {
        var instructionTextbox = document.getElementById('instruction-textbox');
        if (instructionTextbox) {
            var textarea = instructionTextbox.querySelector('textarea');
            if (textarea) {
                textarea.addEventListener('keydown', onKeyPress);
            }
        }
    });

    // Gradio can be slow to load, so we'll also use a MutationObserver
    // to make sure we attach the event listener even if the element appears later.
    const observer = new MutationObserver((mutations, obs) => {
        const instructionTextbox = document.getElementById('instruction-textbox');
        if (instructionTextbox) {
            const textarea = instructionTextbox.querySelector('textarea');
            if (textarea) {
                textarea.addEventListener('keydown', onKeyPress);
                obs.disconnect(); // Stop observing once we've found it
            }
        }
    });
    observer.observe(document.body, {
        childList: true,
        subtree: true
    });
    </script>
    """
    gr.HTML(js_code)


# This will launch the app when deployed on Hugging Face Spaces
demo.launch()