import gradio as gr
from poe_api_wrapper import PoeApi
from typing import Generator

def init_client(p_b: str, p_lat: str) -> PoeApi:
    return PoeApi(token={"p-b": p_b, "p-lat": p_lat})

custom_css = """
#component-0 {
    background: linear-gradient(135deg, #8B5CF6 0%, #FCD34D 100%);
    padding: 20px;
    border-radius: 15px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}
.gradio-container {
    background: linear-gradient(135deg, #8B5CF6 0%, #FCD34D 100%);
}
.chat-message {
    padding: 15px;
    margin: 5px;
    border-radius: 10px;
    background-color: rgba(255, 255, 255, 0.9);
}
.assistant-message {
    border-left: 4px solid #8B5CF6;
}
.user-message {
    border-left: 4px solid #FCD34D;
}
"""

def chat_stream(client: PoeApi, message: str, bot: str, history: list) -> Generator:
    response = ""
    for chunk in client.send_message(bot, message):
        response += chunk["response"]
        yield response

def create_app():
    with gr.Blocks(css=custom_css) as app:
        gr.Markdown("# 🤖 Poe Chat Interface")
        
        with gr.Row():
            with gr.Column():
                p_b = gr.Textbox(label="P-B Token", type="password")
                p_lat = gr.Textbox(label="P-LAT Token", type="password")
                bot_select = gr.Dropdown(
                    choices=["Claude-3-Opus", "Claude-3-Sonnet", "GPT-4", "GPT-3.5-Turbo", "Claude-2", "Google-PaLM"],
                    label="Select Bot",
                    value="Claude-3.7-Sonnet"
                )
                connect_btn = gr.Button("Connect", variant="primary")

        chatbot = gr.Chatbot(
            [],
            elem_classes=["chat-message", "assistant-message", "user-message"],
            height=500
        )
        msg = gr.Textbox(
            label="Type your message here...",
            placeholder="Enter your message and press enter...",
            lines=3
        )
        clear = gr.Button("Clear Chat")

        client = {"instance": None}

        def connect(p_b_val: str, p_lat_val: str) -> str:
            try:
                client["instance"] = init_client(p_b_val, p_lat_val)
                return "✅ Successfully connected!"
            except Exception as e:
                # Log the error for debugging
                print(f"Connection error: {e}")
                return f"❌ Connection failed: {str(e)}"

        def respond(message: str, chat_history: list) -> tuple:
            if not client["instance"]:
                return chat_history + [[message, "Please connect to Poe first!"]]
            
            bot = bot_select.value.lower().replace("-", "").replace(" ", "")
            history = chat_history + [[message, ""]]
            
            for response in chat_stream(client["instance"], message, bot, history):
                history[-1][1] = response
                yield history

        def clear_chat() -> tuple:
            return [], []

        connect_btn.click(connect, [p_b, p_lat], gr.Textbox(label="Connection Status"))
        msg.submit(respond, [msg, chatbot], [chatbot])
        clear.click(clear_chat, None, [chatbot, msg])

    return app

demo = create_app()

if __name__ == "__main__":
    demo.launch()
else:
    app = create_app()