import os
import torch
import streamlit as st
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig
)
from peft import PeftModel

# Offload directory for CPU inference
os.makedirs("offload", exist_ok=True)

# Load base model + quantization config
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,
    llm_int8_enable_fp32_cpu_offload=True
)

tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    offload_folder="offload"
)

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "lora_adapter")

# Evaluation mode
model.eval()

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

### Instruction:
{instruction}

### Response:
"""

# 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=True,
        temperature=0.7,
        top_p=0.95,
        repetition_penalty=1.2
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("### Response:")[-1].strip()

# Streamlit UI
st.title("🐍 Python Tutor Chatbot")
st.markdown("Ask me Python programming questions!")

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

if st.button("Answer"):
    if user_input.strip():
        with st.spinner("Thinking..."):
            answer = chat(user_input)
        st.markdown("**Answer:**")
        st.write(answer)
    else:
        st.warning("Please enter a question.")