import json
from typing import Optional

from fastapi import Body, HTTPException
from fastapi.responses import JSONResponse
from app.comfy.comfy_queue import queue_manager
from app.comfy.comfy_workflow import ComfyWorkflow
from app.database.api import get_workflow
from ._base import runx_router

RUNX_NAME = "model_with_multi_dress"
RUNX_WORKFLOW_NAME = "测试"


@runx_router.post(f"/{RUNX_NAME}")
async def model_with_multi_dress(
    data: dict = Body(...),
    workflow_version: Optional[str] = None,
):
    """
    一个模特，穿多件不同的衣服
    """
    workflow_name = RUNX_WORKFLOW_NAME

    # 获取工作流定义
    workflow_data = await get_workflow(workflow_name, workflow_version)
    if not workflow_data:
        detail = (
            f"工作流 '{workflow_name}'"
            + (f" 带版本 '{workflow_version}'" if workflow_version else " (最新版)")
            + " 未找到。"
        )
        raise HTTPException(status_code=404, detail=detail)

    workflow = json.loads(workflow_data["workflow_json"])
    flow = ComfyWorkflow(workflow_name, workflow)

    # 将请求拆分为多个请求
    batch_data = _convert(data)

    # 提交到队列
    workflow_run_ids: list[str] = []
    for item in batch_data:
        workflow_run_id = await queue_manager.add_task(workflow=flow, request_data=item)
        workflow_run_ids.append(workflow_run_id)

    return JSONResponse(
        content={
            "workflow_run_ids": workflow_run_ids,
            "status": "queued",
            "message": "工作流已提交到队列，正在等待执行",
        },
        status_code=202,
    )


@runx_router.get(f"/{RUNX_NAME}/json_schema")
async def model_with_multi_dress_json_schema():
    """
    获取工作流定义
    """
    return {
        "type": "object",
        "properties": {
            "模特图": {
                "type": "object",
                "properties": {
                    "image": {"type": "string"},
                },
            },
            "模特描述": {
                "type": "object",
                "properties": {
                    "value": {"type": "string"},
                },
            },
            "穿搭图": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "image": {"type": "string"},
                    },
                },
            },
        },
        "required": ["模特图", "模特描述", "穿搭图"],
    }


def _convert(data: dict) -> list[dict]:
    """
    数据格式
    {
        "模特图": {"image" : "xxx" },
        "模特描述" : { "value" : "xxx" },
        "穿搭图": [
            {"image" : "xxx" },
            {"image" : "xxx" }
        ]
    }
    转换为
    [
        {
            "模特图": {"image" : "xxx" },
            "模特描述" : { "value" : "xxx" },
            "穿搭图": { "image" : "xxx" }
        }
    ]
    """
    result: list[dict] = []

    # 获取基础信息
    model_image = data.get("模特图", {})
    model_description = data.get("模特描述", {})
    dress_images = data.get("穿搭图", [])

    # 为每个穿搭图创建一个记录
    for dress_image in dress_images:
        record = {
            "模特图": model_image,
            "模特描述": model_description,
            "穿搭图": dress_image,
        }
        result.append(record)

    return result