import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import streamlit as st

# Load base model and tokenizer
base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(base_model)

# Load model in CPU mode
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    torch_dtype=torch.float32,
    device_map="cpu"
)

# Load the LoRA adapter
model = PeftModel.from_pretrained(model, "lora_adapter", device_map="cpu")
model.eval()

# Format prompt with example
def format_prompt(instruction):
    return f"""### SYSTEM:
You are a helpful and expert Python programming tutor.
You only answer questions related to Python programming.
If the question is unrelated to Python, say:
"Sorry, I can only answer Python-related questions."

### USER:
What is the difference between a list and a tuple in Python?

### ASSISTANT:
In Python, both lists and tuples are used to store collections of items, but they have key differences. Lists are mutable (can be changed), whereas tuples are immutable (cannot be changed). Lists use square brackets [], and tuples use parentheses (). Tuples are generally faster and use less memory.

### USER:
{instruction}

### ASSISTANT:
"""

# Chat function
def chat(instruction):
    prompt = format_prompt(instruction)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        temperature=0.0,
        top_p=1.0,
        repetition_penalty=1.1
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print("🧠 Raw Model Output:", repr(response))  # Debug log

    # Extract response cleanly
    if "### ASSISTANT:" in response:
        return response.split("### ASSISTANT:")[-1].strip()
    else:
        return response.strip()  # fallback

# Streamlit UI
st.set_page_config(page_title="🐍 Python Tutor Chatbot")
st.title("🐍 Python Tutor Chatbot")
st.write("Ask me Python programming questions!")

user_input = st.text_area("Your question:")

if st.button("Get Answer") and user_input.strip():
    with st.spinner("Thinking..."):
        response = chat(user_input)
        st.markdown("**Answer:**")
        st.write(response)

        # Optional debug output
        # st.write("**Raw model output:**")
        # st.write(repr(response))