import gradio as gr
import httpx
import uuid
import json
from collections import OrderedDict

API_URL = "https://oluinioluwa814-telecom.hf.space/stream"

# SESSION STATE (UI ONLY)
sessions = OrderedDict()
MAX_SESSIONS = 50

def get_session_id(current_id):
    if current_id in sessions:
        sessions.move_to_end(current_id)
        return current_id
    if len(sessions) >= MAX_SESSIONS:
        sessions.popitem(last=False)
    new_id = str(uuid.uuid4())[:8]
    sessions[new_id] = []
    return new_id
    
# STREAM CHAT
async def chat(message, history, session_id):
    if not message.strip():
        yield history, session_id, gr.update()
        return
    active_id = get_session_id(session_id)
    if history is None:
        history = []
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": ""})
    yield history, active_id, gr.update(
        choices=list(sessions.keys()),
        value=active_id
    )
    full_response = ""
    try:
        async with httpx.AsyncClient(timeout=None) as client:
            async with client.stream(
                "POST",
                API_URL,
                json={"query": message, "session": active_id},
            ) as response:
                if response.status_code != 200:
                    raise Exception(f"Bad response: {response.status_code}")                    
                async for line in response.aiter_lines():
                    if not line.strip():
                        continue
                    try:
                        data = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    if "error" in data:
                        raise Exception(data["error"])
                    chunk = data.get("text", "")
                    full_response += chunk
                    history[-1]["content"] = full_response
                    yield history, active_id, gr.update()

    except Exception as e:
        print("STREAM ERROR:", e)
        history[-1]["content"] = "⚠️ Server error. Try again."
        yield history, active_id, gr.update()
        
# SESSION SWITCH
def switch_session(new_id):
    if not new_id:
        return [], None
    return sessions.get(new_id, []), new_id

def clear_session(s):
    if s:
        sessions[s] = []
    return [], s

# UI
with gr.Blocks() as demo:
    gr.Markdown("# 📡 Telecom AI Assistant")
    session_state = gr.State()
    with gr.Row():
        with gr.Column(scale=1):
            session_drop = gr.Dropdown(
                label="Active Chats",
                choices=[],
                interactive=True,
            )
            new_btn = gr.Button("➕ New Chat", variant="secondary")
            clear_btn = gr.Button("🗑 Clear History", variant="stop")
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(height=550)
            with gr.Row():
                msg = gr.Textbox(
                    placeholder="Ask me anything...",
                    show_label=False,
                    scale=9,
                )
                send_btn = gr.Button("➤", scale=1, variant="primary")

    # shared function
    chat_args = {
        "fn": chat,
        "inputs": [msg, chatbot, session_state],
        "outputs": [chatbot, session_state, session_drop],
    }
    msg.submit(**chat_args).then(lambda: "", None, msg)
    send_btn.click(**chat_args).then(lambda: "", None, msg)
    new_btn.click(
        lambda: ([], None, gr.update(value=None)),
        None,
        [chatbot, session_state, session_drop],
    )
    clear_btn.click(
        clear_session,
        session_state,
        [chatbot, session_state]
    )
    session_drop.change(
        switch_session,
        session_drop,
        [chatbot, session_state]
    )
if __name__ == "__main__":
    demo.queue().launch(
        ssr_mode=False,
        theme=gr.themes.Default(primary_hue="blue")
    )