import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from threading import Thread
import time

# Load small and fast model for Python code debugging
MODEL_NAME = "Salesforce/codet5-small"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def debug_python(code, progress=gr.State(0)):
    if not code.strip():
        return "❌ Please enter Python code.", 0

    prompt = f"Fix this Python code:\n{code}\nCorrected code:\n"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)

    def update_progress():
        for i in range(1, 101):
            time.sleep(0.005)  # Faster loading feel
            progress.value = i

    Thread(target=update_progress).start()

    outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.2, do_sample=False)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    progress.value = 100
    return response.strip(), 100

with gr.Blocks(css=".gradio-container {background-color: #cbedec;}") as app:
    gr.Markdown("## 🐍 Python Code Debugger (Fast Mode)")
    gr.Markdown("Just paste your Python code. The AI will fix any issues it detects.")

    code_box = gr.Textbox(label="📝 Your Python Code", lines=12)
    progress_bar = gr.Slider(minimum=0, maximum=100, value=0, interactive=False, label="⏳ Progress (%)")
    output = gr.Code(label="✅ Suggested Fix")
    fix_btn = gr.Button("🛠️ Debug Code")

    fix_btn.click(fn=debug_python, inputs=[code_box], outputs=[output, progress_bar])

app.launch()