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

MODEL_ID = "Qwen/Qwen2.5-Coder-1.5B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
)

SYSTEM_PROMPT = """You are an expert code reviewer. Analyze the provided code and give structured feedback in this exact format:

## 🐛 Bugs & Errors
List any bugs, logic errors, or runtime issues. If none, say "None found."

## 🔒 Security Issues
List any security vulnerabilities (SQL injection, XSS, hardcoded secrets, etc.). If none, say "None found."

## 🎨 Code Style & Readability
Comment on naming conventions, formatting, comments, and clarity.

## ⚡ Performance
Identify any inefficiencies or bottlenecks.

## ✅ Suggestions
Top 3 actionable improvements the developer should make.

Be specific, concise, and constructive. Reference line numbers where possible."""


def build_prompt(code: str, language: str) -> str:
    return f"""<s>[INST] <<SYS>>
{SYSTEM_PROMPT}
<</SYS>>

Please review the following {language} code:

```{language.lower()}
{code}
```
[/INST]"""


def review_code(code: str, language: str, max_new_tokens: int = 800):
    try:
        if not code.strip():
            return "⚠️ Please paste some code to review."
    
        prompt = build_prompt(code, language)
        inputs = tokenizer(prompt, return_tensors="pt")
    
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
                temperature=0.2,
                do_sample=True,
                top_p=0.95,
                repetition_penalty=1.1,
                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.eos_token_id,
                use_cache=True
            )
    
        # Decode only the newly generated tokens (skip the prompt)
        response = tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True
        )
        return response.strip() or "⚠️ Model returned an empty response."
    except Exception as e:
        return f"❌ Error:\n\n{str(e)}"

# --- Gradio UI ---
with gr.Blocks(title="AI Code Reviewer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔍 AI Code Reviewer\nPowered by CodeLlama-7B · Paste your code and get structured feedback instantly.")

    with gr.Row():
        with gr.Column(scale=2):
            language = gr.Dropdown(
                choices=["Python", "JavaScript", "TypeScript", "Java", "C++", "Go", "Rust", "SQL"],
                value="Python",
                label="Language"
            )
            code_input = gr.Code(
                language="python",
                label="Your Code",
                lines=20,
                value="# Paste your code here...",
            )
            submit_btn = gr.Button("Review Code ↗", variant="primary")

        with gr.Column(scale=2):
            output = gr.Markdown(label="Review", value="*Review will appear here...*")

    submit_btn.click(
        fn=review_code,
        inputs=[code_input, language],
        outputs=output,
        show_progress="full"
    )

    gr.Examples(
        examples=[
            ["def divide(a, b):\n    return a / b\n\nresult = divide(10, 0)\nprint(result)", "Python"],
            ["SELECT * FROM users WHERE name = '" + "' + username + '", "SQL"],
        ],
        inputs=[code_input, language],
        label="Try an example"
    )

demo.queue(max_size=10).launch()