from typing import Tuple, Optional, Any, List

import modal
from loguru import logger
from modal.call_graph import InputStatus
from pydantic import BaseModel, Field

from BowongModalFunctions.models.web_model import TaskStatus, SentryTransactionInfo, ErrorCode


class ModalTaskInfo(BaseModel):
    function_name: str = Field(description="任务名字")
    status: TaskStatus = Field(description="任务运行状态")
    error_code: ErrorCode = Field(default=ErrorCode.SUCCESS, description="错误代码")
    error_reason: Optional[str] = Field(default=None, description="错误原因")
    results: Optional[List[Any]] = Field(default=None, description="任务结果")
    transaction: Optional[SentryTransactionInfo] = Field(default=None, description="Sentry跟踪信息")


class ModalUtils:
    @staticmethod
    async def get_modal_task_status(task_id: str) -> ModalTaskInfo:
        """
        使用modal任务id查看任务运行状态和结果
        :param task_id: modal 任务id
        :return: (TaskStaus 任务运行状态, function_name 任务名字, errorCode 错误代码, errorReason 错误原因, results 任务结果, sentryTransactionInfo Sentry跟踪信息)
        """
        task = None
        try:
            fn_task = modal.FunctionCall.from_id(task_id)

            call_graph = fn_task.get_call_graph()
            root = call_graph[0]
            if len(root.children) == 0:
                return ModalTaskInfo(function_name="unknown", error_code=ErrorCode.NOT_FOUND,
                                     status=TaskStatus.expired, error_reason="NOT_FOUND")
            task = root.children[0]
            if not task.function_call_id == task_id:
                return ModalTaskInfo(function_name=task.function_name, status=TaskStatus.expired,
                                     error_code=ErrorCode.NOT_FOUND, error_reason="NOT_FOUND")
            logger.info(f"Task {task.function_name} function call status: {task.status}")
            match task.status:
                case InputStatus.PENDING:
                    return ModalTaskInfo(function_name=task.function_name, status=TaskStatus.running)
                case InputStatus.SUCCESS:
                    try:
                        result, sentry_trace = fn_task.get(timeout=2)
                        if isinstance(result, List):
                            results = result
                        else:
                            results = [result]
                        return ModalTaskInfo(function_name=task.function_name, status=TaskStatus.success,
                                             error_code=ErrorCode.SUCCESS, results=results,
                                             transaction=sentry_trace)
                    except modal.exception.OutputExpiredError:
                        return ModalTaskInfo(function_name=task.function_name, status=TaskStatus.expired,
                                             error_code=ErrorCode.NOT_FOUND, error_reason="EXPIRED")
                    except TimeoutError:
                        return ModalTaskInfo(function_name=task.function_name, status=TaskStatus.running)
                    except Exception as e:
                        logger.exception(e)
                        return ModalTaskInfo(function_name=task.function_name, status=TaskStatus.running,
                                             error_code=ErrorCode.SYSTEM_ERROR, error_reason="FAILURE")
                case _:
                    error = 'FAILURE'
                    match task.status:
                        case InputStatus.INIT_FAILURE:
                            error = 'INIT_FAILURE'
                        case InputStatus.TERMINATED:
                            error = 'TERMINATED'
                        case InputStatus.TIMEOUT:
                            error = 'TIMEOUT'
                    return ModalTaskInfo(function_name=task.function_name, status=TaskStatus.failed,
                                         error_code=ErrorCode.SYSTEM_ERROR, error_reason=error)
        except Exception as e:
            logger.exception(e)
            return ModalTaskInfo(function_name=task.function_name if task else "unknown", status=TaskStatus.failed,
                                 error_code=ErrorCode.NOT_FOUND, error_reason="NOT_FOUND")

    @staticmethod
    async def cancel_modal_task(task_id: str):
        try:
            fn_task = modal.FunctionCall.from_id(task_id)
            fn_task.cancel()
        except Exception as e:
            logger.exception(e)
            return TaskStatus.failed, ErrorCode.NOT_FOUND.value, "NOT_FOUND", None, None
