import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ── Config ──────────────────────────────────────────────────────────────
BASE_MODEL    = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
ADAPTER_REPO = "TRUFELLINI/qwen2.5-coder-0.5b-bugfix"  

SYSTEM_PROMPT = """You are a Python bug fixing assistant.
Rules:
- Make the minimal change possible
- Do not rewrite unrelated code
- Return only fixed code
- Preserve original behavior unless required"""

# ── Load model (runs once when Space starts) ─────────────────────────────
print("Loading base model...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float32,  # CPU — float32
    device_map="cpu",
)

print("Loading LoRA adapter...")
model = PeftModel.from_pretrained(model, ADAPTER_REPO)
model.eval()
print("Model ready.")

# ── Inference ────────────────────────────────────────────────────────────
def fix_bug(buggy_code, max_new_tokens=400):
    if not buggy_code.strip():
        return "Paste some buggy Python code above."

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": f"Fix the bug:\n{buggy_code}"}
    ]
    text   = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt")

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens = max_new_tokens,
            temperature    = 0.1,
            do_sample      = True,
            pad_token_id   = tokenizer.eos_token_id,
        )
    new_tokens = output[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True)

# ── UI ───────────────────────────────────────────────────────────────────
EXAMPLES = [
    ["def calculate_average(numbers):\n    total = 0\n    for n in numbers:\n        total += n\n    return total / len(numbers)"],
    ["def add_item(item, items=[]):\n    items.append(item)\n    return items"],
    ["def sum_list(nums):\n    total = 0\n    for i in range(len(nums)+1):\n        total += nums[i]\n    return total"],
    ["def first_element(lst):\n    return lst[1]"],
]

demo = gr.Interface(
    fn          = fix_bug,
    inputs      = gr.Code(
                    label       = "Paste buggy Python code here",
                    language    = "python",
                    lines       = 14,
                  ),
    outputs     = gr.Code(
                    label       = "Fixed code",
                    language    = "python",
                  ),
    title       = "🐛 Python Bug Fixer",
    description = """**Qwen2.5-Coder 0.5B fine-tuned on real buggy→fixed Python pairs.**  
Makes minimal, correct fixes — not rewrites. Trained using QLoRA on [alexjercan/bugnet](https://huggingface.co/datasets/alexjercan/bugnet).  
[[Model]](https://huggingface.co/your-username/qwen2.5-coder-0.5b-bugfix) · [[GitHub]](https://github.com/your-username/python-bugfix-llm)""",
    examples    = EXAMPLES,
    theme       = gr.themes.Soft(),
    cache_examples = False,
)

demo.launch()