import os
import streamlit as st
from transformers import pipeline
import datetime
import pyttsx3
import tempfile
import speech_recognition as sr

st.set_page_config(page_title="Winner Chat", layout="wide", initial_sidebar_state="expanded")

# Sidebar: Model selector + Dark mode
st.sidebar.title("⚙️ Settings")
model_name = st.sidebar.selectbox("Choose model", ["Arni1ntares/Winner", "gpt2", "distilgpt2"])
dark_mode = st.sidebar.toggle("🌙 Dark Mode")
enable_tts = st.sidebar.toggle("🔊 Enable Voice Output")
enable_mic = st.sidebar.toggle("🎙️ Enable Voice Input")

# Voice recognizer
recognizer = sr.Recognizer()

@st.cache_resource(show_spinner="Loading model...")
def load_model(model_name):
    return pipeline("text-generation", model=model_name)

pipe = load_model(model_name)

# Init TTS
tts = pyttsx3.init()
tts.setProperty('rate', 165)

# Custom CSS for message bubbles and dark mode
def add_css():
    css = """
    <style>
    .user-msg {
        background-color: #DCF8C6;
        padding: 0.5em 1em;
        border-radius: 10px;
        margin: 0.3em 0;
    }
    .bot-msg {
        background-color: #F1F0F0;
        padding: 0.5em 1em;
        border-radius: 10px;
        margin: 0.3em 0;
    }
    .avatar {
        width: 30px; height: 30px; border-radius: 50%;
        display: inline-block; vertical-align: middle;
    }
    .msg-container {
        display: flex;
        align-items: flex-start;
        gap: 10px;
    }
    .bubble-container {
        max-width: 80%;
    }
    </style>
    """
    if dark_mode:
        css += """
        <style>
        body {
            background-color: #0e1117;
            color: white;
        }
        .bot-msg {
            background-color: #2c2f36;
            color: white;
        }
        .user-msg {
            background-color: #005c4b;
            color: white;
        }
        </style>
        """
    st.markdown(css, unsafe_allow_html=True)

add_css()

st.title("🤖 Winner Chat – Enhanced UI")

# Init chat state
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat history
for msg in st.session_state.messages:
    avatar = "🧑" if msg["role"] == "user" else "🤖"
    bubble_class = "user-msg" if msg["role"] == "user" else "bot-msg"
    with st.container():
        st.markdown(f"""
        <div class="msg-container">
            <div class="avatar">{avatar}</div>
            <div class="bubble-container"><div class="{bubble_class}">{msg['content']}</div></div>
        </div>
        """, unsafe_allow_html=True)

# Voice input
user_input = ""
if enable_mic:
    st.subheader("🎙️ Voice Input")
    mic_button = st.button("Start Listening")
    if mic_button:
        with sr.Microphone() as source:
            st.info("Listening...")
            audio = recognizer.listen(source)
            try:
                user_input = recognizer.recognize_google(audio)
                st.success(f"You said: {user_input}")
            except Exception as e:
                st.error(f"Error: {e}")
else:
    user_input = st.chat_input("Type your message...")

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

    with st.spinner("Thinking..."):
        full_prompt = "\n".join(
            [f"{m['role']}: {m['content']}" for m in st.session_state.messages]
        )
        result = pipe(full_prompt, max_new_tokens=200, temperature=0.7)[0]["generated_text"]
        response = result.replace(full_prompt, "").strip()

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

    if enable_tts:
        tts.say(response)
        tts.runAndWait()

    st.experimental_rerun()

# Export conversation
st.sidebar.markdown("---")
if st.sidebar.button("📥 Export Chat"):
    now = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    fname = f"chat_{now}.txt"
    text = "\n".join([f"{m['role']}: {m['content']}" for m in st.session_state.messages])
    with tempfile.NamedTemporaryFile(delete=False, mode="w", suffix=".txt") as f:
        f.write(text)
        st.sidebar.download_button("Download Chat", text, file_name=fname, mime="text/plain")

st.markdown("---")
st.caption("🛠 Built with Streamlit, Hugging Face, and 💛 for AI.")