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

# ==============================
# PAGE CONFIG
# ==============================
st.set_page_config(page_title="💻 AI Coding Assistant", layout="wide")
st.title("💻 AI Coding Assistant")

# ==============================
# LOAD MODEL (LIGHTWEIGHT)
# ==============================
@st.cache_resource
def load_model():
    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float32
    )

    model.eval()
    return tokenizer, model

with st.spinner("🔄 Loading model..."):
    tokenizer, model = load_model()

st.success("✅ Ready")

# ==============================
# SESSION STATE (CHAT HISTORY)
# ==============================
if "messages" not in st.session_state:
    st.session_state.messages = []

# ==============================
# CHAT DISPLAY
# ==============================
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        if msg["role"] == "assistant":
            st.code(msg["content"])
        else:
            st.markdown(msg["content"])

# ==============================
# GENERATE RESPONSE
# ==============================
def generate_response(user_input):

    # Build conversation prompt
    conversation = ""
    for msg in st.session_state.messages:
        role = "User" if msg["role"] == "user" else "Assistant"
        conversation += f"{role}: {msg['content']}\n"

    conversation += f"User: {user_input}\nAssistant:"

    inputs = tokenizer(conversation, return_tensors="pt", truncation=True)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.3,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id
        )

    result = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extract only assistant reply
    if "Assistant:" in result:
        result = result.split("Assistant:")[-1]

    return result.strip()

# ==============================
# CHAT INPUT
# ==============================
user_input = st.chat_input("Ask your coding question...")

if user_input:
    # Add user message
    st.session_state.messages.append({"role": "user", "content": user_input})

    with st.chat_message("user"):
        st.markdown(user_input)

    # Generate response
    with st.spinner("💡 Thinking..."):
        response = generate_response(user_input)

    # Add assistant message
    st.session_state.messages.append({"role": "assistant", "content": response})

    with st.chat_message("assistant"):
        st.code(response)