import gradio as gr
import os
from bedrock_client import get_anthropic_client, claude_stream_response

# Create Claude client via the official SDK
client = get_anthropic_client()

# Streaming chat handler
async def chat(user_message, history):
    messages = history.copy()
    messages.append({"role": "user", "content": user_message})

    full_response = ""

    def run_generator():
        for chunk in claude_stream_response(messages, client):
            yield chunk

    for chunk in run_generator():
        full_response += chunk
        yield messages + [{"role": "assistant", "content": full_response}]

    # Append full assistant response to history
    messages.append({"role": "assistant", "content": full_response})
    yield messages

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 Claude v2.1 Chatbot via Amazon Bedrock (`anthropic[bedrock]`)")
    chatbot = gr.Chatbot(label="Claude v2.1", type="messages")
    msg = gr.Textbox(placeholder="Ask Claude anything...", show_label=False)
    state = gr.State([])  # chat history
    clear = gr.Button("🧹 Clear Chat")

    def clear_chat():
        return [], []

    msg.submit(chat, [msg, state], chatbot, queue=True)
    msg.submit(lambda: "", None, msg)  # clear input
    clear.click(clear_chat, None, [chatbot, state])

demo.queue().launch()