from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load Model
model_id = "bigcode/starcoder2-3b"
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_4bit=True,
    torch_dtype=torch.float16,
    device_map="cpu"
)

# ✅ CodeBuddy System Prompt
system_prompt = """
You are **CodeBuddy**, a friendly and skilled developer assistant inside a code editor.

Your job:
- Explain code clearly
- Suggest improvements
- Fix bugs when asked
- Teach best practices
- Keep your answers helpful, accurate, and detailed
- Always speak in a friendly, developer-to-developer tone

Focus languages for now:
- JavaScript
- HTML
- CSS

Important rules:
1. When explaining code, break it down step-by-step.
2. When giving a fix, ALWAYS show corrected code in a formatted code block.
3. When suggesting improvements, explain *why* the change is beneficial.
4. Avoid unnecessary technical jargon unless the user is clearly advanced.
5. If the user provides broken code, help debug it — don’t just rewrite everything unless needed.
6. If unsure about something, say so briefly but try to reason your way through.

Tone:
- Friendly dev talk
- Helpful, not robotic
- Encouraging, not formal
"""

# 🔥 Example user input
user_input = """
Explain this code:

for i in range(5):
    print(i)
"""

# 🧠 Combine System Prompt + User Input
prompt = system_prompt + "\n\nUser:\n" + user_input + "\n\nCodeBuddy:"

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

output = model.generate(
    **inputs,
    max_new_tokens=200,
    temperature=0.4,
    top_p=0.9,
)

print(tokenizer.decode(output[0], skip_special_tokens=True))