import streamlit as st
import requests
import os

# === Constants ===
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY") or st.secrets.get("OPENROUTER_API_KEY", "")
MODEL_NAME = "deepseek/deepseek-r1-0528:free"
API_URL = "https://openrouter.ai/api/v1/chat/completions"

# === Streamlit UI Setup ===
st.title("💬 Whatsapp like Chat with DeepSeek-R1")

# Custom WhatsApp-style CSS
st.markdown("""
<style>
    .stChatMessage { background-color: transparent !important; padding: 0 !important; width: 100%; }
    .user-container, .assistant-container { display: flex; flex-direction: column; margin-bottom: 12px; }
    .user-container { align-items: flex-end; }
    .assistant-container { align-items: flex-start; }
    .user-message, .assistant-message {
        padding: 8px 12px; border-radius: 15px; max-width: 80%; margin-top: 4px;
    }
    .user-message { background-color: #DCF8C6; border-radius: 15px 15px 0 15px; }
    .assistant-message { background-color: #FFF; border: 1px solid #ECECEC; border-radius: 15px 15px 15px 0; }
    .sender-label { font-weight: bold; font-size: 0.8em; color: #555; }
    .user-label { margin-right: 8px; }
    .assistant-label { margin-left: 8px; }
    [data-testid="stChatMessage"] { background-color: transparent !important; }
    [data-testid="stChatMessage"] > div:first-child { display: none !important; }
</style>
""", unsafe_allow_html=True)

# === Chat History ===
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": "You are a helpful assistant."}]

# === Display Chat History ===
for msg in st.session_state.messages:
    if msg["role"] == "system":
        continue
    with st.chat_message(msg["role"]):
        content_html = f"""
        <div class="{msg['role']}-container">
            <div class="{msg['role']}-label sender-label">{'You' if msg['role']=='user' else 'Bot'}:</div>
            <div class="{msg['role']}-message">{msg['content']}</div>
        </div>
        """
        st.markdown(content_html, unsafe_allow_html=True)

# === User Input ===
if prompt := st.chat_input("Type your message..."):
    # Add user message
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(f"""
        <div class="user-container">
            <div class="user-label sender-label">Me:</div>
            <div class="user-message">{prompt}</div>
        </div>
        """, unsafe_allow_html=True)

    with st.chat_message("assistant"):
        placeholder = st.empty()
        placeholder.markdown("""
        <div class="assistant-container">
            <div class="assistant-label sender-label">Bot:</div>
            <div class="assistant-message">▌</div>
        </div>
        """, unsafe_allow_html=True)

        # === OpenRouter API Call ===
        try:
            headers = {
                "Authorization": f"Bearer {OPENROUTER_API_KEY}",
                "Content-Type": "application/json",
                "HTTP-Referer": "https://your-app-name.streamlit.app/",  # required for free tier
                "X-Title": "Streamlit WhatsApp Chatbot"
            }
            payload = {
                "model": MODEL_NAME,
                "messages": st.session_state.messages,
                "stream": False
            }

            response = requests.post(API_URL, headers=headers, json=payload)
            response.raise_for_status()
            full_response = response.json()["choices"][0]["message"]["content"]

        except Exception as e:
            full_response = f"Error: {e}"

        # Display and store assistant message
        placeholder.markdown(f"""
        <div class="assistant-container">
            <div class="assistant-label sender-label">Bot:</div>
            <div class="assistant-message">{full_response}</div>
        </div>
        """, unsafe_allow_html=True)

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