import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load CodeT5+ model and tokenizer
model_id = "Salesforce/codet5p-770m"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(model_id, torch_dtype=torch.float32)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Prompt templates for supported languages
language_prompts = {
    "Python": "Fix the following Python code:\n",
    "C": "Fix the following C code:\n",
    "C++": "Fix the following C++ code:\n",
    "JavaScript": "Fix the following JavaScript code:\n"
}

# Debug function
def eternos_debugger(code, error, language):
    if not code.strip():
        return "❌ Please enter some code to debug."
    prompt = language_prompts[language] + code + "\nError:\n" + error + "\nFixed Code:\n"

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.3,
            top_p=0.9,
            do_sample=True
        )
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result.strip()

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🛠️ Eternos — AI Code Debugger")
    gr.Markdown("Supports Python, C, C++, JavaScript — Powered by CodeT5+")

    with gr.Row():
        code_input = gr.Textbox(label="👨‍💻 Your Code", lines=14, placeholder="Paste your buggy code here...")
        error_input = gr.Textbox(label="⚠️ Error Message (optional)", lines=4)

    language_input = gr.Dropdown(["Python", "C", "C++", "JavaScript"], label="🌐 Language", value="Python")
    output_code = gr.Code(label="✅ Suggested Fix")
    run_btn = gr.Button("🔍 Debug Code")

    run_btn.click(fn=eternos_debugger, inputs=[code_input, error_input, language_input], outputs=output_code)

demo.launch()