import json
import gradio as gr

from coze import run_workflow, call_plugin


async def generate_image(query):
    parameters = {
        "text_prompts": [{"text": query}],
        # Stable Diffusion XL v1.0
        "model_type": 2,
        "width": 1344,
        "height": 768,
    }
    async for message in call_plugin(parameters, "Stable Diffusion:textToImage"):
        event = message.get("event")
        data = message.get("data")

        if event == "conversation.message.completed" and data:
            try:
                data = json.loads(data)  # 尝试解析 JSON 数据
                response_type = data.get("type")
                if response_type == "tool_response":
                    content = json.loads(data.get("content"))
                    image_url = content.get("data").get("images")[0]
                    if image_url:
                        yield f'<img src="{image_url})" />'
            except json.JSONDecodeError:
                yield "Error: Failed to decode JSON data."


async def create_full_story(query, history=None):
    parameters = {"query": query}
    full_response = ""
    async for message in run_workflow(parameters, "self:create_full_story"):
        event = message.get("event")
        if event == "conversation.message.delta":
            data = json.loads(message.get("data"))
            content = data.get("content")
            if content:
                full_response += content
                yield full_response


with gr.Blocks() as demo:
    with gr.Tab("Call Coze Workflow"):
        gr.ChatInterface(create_full_story, fill_height=True)

    with gr.Tab("Call Coze Plugin"):
        prompt = gr.Textbox(
            label="输入描述", value="A beautiful sunset over the ocean."
        )
        output = gr.HTML()
        generate_button = gr.Button("生成图片")
        generate_button.click(fn=generate_image, inputs=prompt, outputs=output)

if __name__ == "__main__":
    demo.launch()