from typing import Annotated, cast

import modal
from fastapi import APIRouter, Depends, Header
from fastapi.responses import Response
from starlette import status

from ..config import WorkerConfig
from ..middleware.authorization import verify_token
from ..models.media_model import MediaSource
from ..utils.ModalUtils import ModalUtils
from ..models.web_model import (SentryTransactionHeader,
                                ModalTaskResponse, SentryTransactionInfo,
                                ComfyTaskRequest, ComfyTaskStatusResponse, TaskStatus)

config = WorkerConfig()

router = APIRouter(prefix="/comfyui")

sentry_header_schema = {
    "x-trace-id": {
        "description": "Sentry Transaction ID",
        "schema": {
            "type": "string",
        }
    },
    "x-baggage": {
        "description": "Sentry Transaction baggage",
        "schema": {
            "type": "string",
        }
    }
}


@router.post("/v1", tags=["ComfyUI"], summary="发起ComfyUIBase任务", description="发起ComfyUIBase任务",
             responses={
                 status.HTTP_200_OK: {
                     "description": "",
                     "headers": sentry_header_schema
                 },
             },
             dependencies=[Depends(verify_token)]
             )
async def comfyui_v1(item: ComfyTaskRequest, headers: Annotated[SentryTransactionHeader, Header()]):
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    cls = modal.Cls.from_name(config.modal_app_name, "ComfyUI", environment_name=config.modal_environment)
    fn_call = cls().api.spawn(item.video_path.path.replace(config.comfyui_s3_input + "/", ""), item.start_time,
                              item.filename_prefix, item.tts_text1,
                              item.tts_text2, item.tts_text3, item.tts_text4, item.anchor_id, item.speed, sentry_trace,
                              item.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.get("/v1/{task_id}", tags=["ComfyUI"], summary="查询ComfyUIBase任务",
            description="查询ComfyUIBase任务",
            responses={
                status.HTTP_200_OK: {
                    "description": "",
                    "headers": sentry_header_schema
                },
            },
            dependencies=[Depends(verify_token)]
            )
async def comfyui_v1_status(task_id: str, response: Response):
    task_info = await ModalUtils.get_modal_task_status(task_id)
    media = None
    if task_info.transaction:
        response.headers["x-trace-id"] = task_info.transaction.x_trace_id
        response.headers["x-baggage"] = task_info.transaction.x_baggage
    if task_info.status == TaskStatus.success:
        task_status = task_info.results[0]["status"]
        if task_status != "success":
            reason = task_info.results[0]["msg"]
        else:
            media = MediaSource.from_str(task_info.results[0]["file_name"])
    return ComfyTaskStatusResponse(taskId=task_id, status=task_status, code=cast(int, task_info.error_code.value),
                                   error=reason, result=media.urn if media else None)


@router.post("/v2", tags=["ComfyUI"], summary="发起ComfyUILatentSync1.5任务",
             description="发起ComfyUILatentSync1.5任务",
             responses={
                 status.HTTP_200_OK: {
                     "description": "",
                     "headers": sentry_header_schema
                 },
             },
             dependencies=[Depends(verify_token)]
             )
async def comfyui_v2(item: ComfyTaskRequest, headers: Annotated[SentryTransactionHeader, Header()]):
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    cls = modal.Cls.from_name(config.modal_app_name, "ComfyUILatentSync15", environment_name=config.modal_environment)
    fn_call = cls().api.spawn(item.video_path.path.replace(config.comfyui_s3_input + "/", ""), item.start_time,
                              item.filename_prefix, item.tts_text1,
                              item.tts_text2, item.tts_text3, item.tts_text4, item.anchor_id, item.speed, sentry_trace,
                              item.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.get("/v2/{task_id}", tags=["ComfyUI"], summary="查询ComfyUILatentSync1.5任务",
            description="查询ComfyUILatentSync1.5任务",
            responses={
                status.HTTP_200_OK: {
                    "description": "",
                    "headers": sentry_header_schema
                },
            },
            dependencies=[Depends(verify_token)]
            )
async def comfyui_v2_status(task_id: str, response: Response):
    task_info = await ModalUtils.get_modal_task_status(task_id)
    media = None

    if task_info.transaction:
        response.headers["x-trace-id"] = task_info.transaction.x_trace_id
        response.headers["x-baggage"] = task_info.transaction.x_baggage
    if task_info.status == "success":
        task_status = task_info.results[0]["status"]
        if task_status != "success":
            reason = task_info.results[0]["msg"]
        else:
            media = MediaSource.from_str(task_info.results[0]["file_name"])
    return ComfyTaskStatusResponse(taskId=task_id, status=task_status, code=cast(int, task_info.error_code.value),
                                   error=reason, result=media.urn if media else "")
