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

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

# ==============================
# LOAD MODEL (CHAT MODEL ✅)
# ==============================
@st.cache_resource
def load_model():
    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # ✅ Best for HF free

    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
# ==============================
if "messages" not in st.session_state:
    st.session_state.messages = []

# ==============================
# CLEAN TEXT
# ==============================
def clean_text(text):
    text = re.sub(r"[^\x00-\x7F]+", "", text).strip()

    # Ensure response completes nicely
    if not text.endswith((".", "!", "?")):
        text += "..."

    return text

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

    prompt = f"""
<|user|>
{user_input}

Give a clear and complete answer.
<|assistant|>
"""

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

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,   # 🔥 prevents cut-off
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id
        )

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

    # Extract assistant response
    if "<|assistant|>" in result:
        result = result.split("<|assistant|>")[-1]

    return clean_text(result)

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

# ==============================
# INPUT BOX
# ==============================
user_input = st.chat_input("Type your message...")

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 response
    st.session_state.messages.append({"role": "assistant", "content": response})

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

# ==============================
# SIDEBAR
# ==============================
st.sidebar.title("⚙️ Options")

if st.sidebar.button("🧹 Clear Chat"):
    st.session_state.messages = []