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

MODEL_NAME = "Qwen/Qwen2.5-Coder-0.5B-Instruct"

# Load tokenizer and model once at startup
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

# Use GPU when available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = model.to(device)
model.eval()


SYSTEM_PROMPT = (
    "You are a concise code reviewer. Given a code snippet, point out bugs, "
    "edge cases, complexity issues, and style improvements. "
    "Be specific and brief."
)


@spaces.GPU
def review_code(code):
    if not code or not code.strip():
        return "Paste some code first."

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT
        },
        {
            "role": "user",
            "content": f"Review this code:\n\n{code}"
        }
    ]

    # Apply the model's chat template
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=4096
    )

    # Move inputs to the same device as the model
    inputs = {
        key: value.to(device)
        for key, value in inputs.items()
    }

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=300,
            do_sample=False,
            temperature=None,
            top_p=None
        )

    # Only decode newly generated tokens
    generated_ids = output_ids[0][inputs["input_ids"].shape[1]:]

    review = tokenizer.decode(
        generated_ids,
        skip_special_tokens=True
    )

    return review.strip()


with gr.Blocks(title="AI Code Reviewer") as demo:

    gr.Markdown(
        "## AI Code Reviewer\n"
        "Paste a code snippet and get feedback on bugs, edge cases, "
        "complexity, and style."
    )

    code_input = gr.Code(
        label="Your code",
        language="python",
        lines=15
    )

    review_btn = gr.Button(
        "Review",
        variant="primary"
    )

    output = gr.Markdown()

    review_btn.click(
        review_code,
        inputs=code_input,
        outputs=output
    )


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