import streamlit as st
import uuid
import json
import asyncio
import websockets

# ================= CONFIG =================
MIDDLEWARE_WS = "wss://partha181098-middleware.hf.space/chat"

st.set_page_config(page_title="Secure AI Chat", layout="centered")

ss = st.session_state

# ================= SESSION =================
def new_session():
    ss.session_id = str(uuid.uuid4())
    ss.messages = []

if "session_id" not in ss:
    new_session()

if "messages" not in ss:
    ss.messages = []

# ================= UI =================
st.title("🔐 Secure AI Chatbot")

if st.button("New Chat"):
    new_session()
    st.rerun()

# Show chat history
for msg in ss.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

user_input = st.chat_input("Ask something...")

# ================= STREAM HANDLER =================
async def ws_stream(user_text, placeholder):
    uri = f"{MIDDLEWARE_WS}?session_id={ss.session_id}"

    async with websockets.connect(uri) as ws:
        await ws.send(user_text)

        full_response = ""

        while True:
            raw = await ws.recv()
            payload = json.loads(raw)

            if payload["type"] == "security_warning":
                st.warning(payload["message"])
                return None

            if payload["type"] == "response":
                full_response += payload["content"]
                placeholder.markdown(full_response)

            if payload["type"] == "complete":
                return full_response

# ================= CHAT FLOW =================
if user_input:
    # Store user message
    ss.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    with st.chat_message("assistant"):
        placeholder = st.empty()

        # IMPORTANT: Streamlit-safe asyncio execution
        response = asyncio.new_event_loop().run_until_complete(
            ws_stream(user_input, placeholder)
        )

        if response:
            ss.messages.append(
                {"role": "assistant", "content": response}
            )