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

# ✅ Use a cache directory that Spaces allows
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"

@st.cache_resource
def load_model():
    model_name = "microsoft/DialoGPT-small"  # switched from -medium
    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/tmp/hf_cache")
    model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir="/tmp/hf_cache")
    return tokenizer, model

tokenizer, model = load_model()

st.set_page_config(page_title="Chatbot 🤖", page_icon="💬", layout="centered")
st.title("🤖 Hugging Face Chatbot (DialoGPT-small)")

if "chat_history_ids" not in st.session_state:
    st.session_state.chat_history_ids = None
if "past_inputs" not in st.session_state:
    st.session_state.past_inputs = []
if "generated_responses" not in st.session_state:
    st.session_state.generated_responses = []

user_input = st.text_input("You: ", "", key="input")

if st.button("Send") and user_input:
    new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")

    if st.session_state.chat_history_ids is not None:
        bot_input_ids = torch.cat([st.session_state.chat_history_ids, new_input_ids], dim=-1)
    else:
        bot_input_ids = new_input_ids

    st.session_state.chat_history_ids = model.generate(
        bot_input_ids,
        max_length=1000,
        pad_token_id=tokenizer.eos_token_id
    )

    bot_output = tokenizer.decode(
        st.session_state.chat_history_ids[:, bot_input_ids.shape[-1]:][0],
        skip_special_tokens=True
    )

    st.session_state.past_inputs.append(user_input)
    st.session_state.generated_responses.append(bot_output)

if st.session_state.generated_responses:
    for i in range(len(st.session_state.generated_responses)):
        st.markdown(f"**You:** {st.session_state.past_inputs[i]}")
        st.markdown(f"**Bot:** {st.session_state.generated_responses[i]}")