import time
from functools import wraps
from logging import Logger
from typing import Any, Callable, Generator

import humanize

from ldnode import logger
from ldnode.calculator.handlers.models import (
    TaskBaseParameters,
    TaskResponse,
)

from .responses import error_response, response


def wrap_handler(
    task_name: str, logger: Logger = logger.calculator, is_generator: bool = False
):
    """
    任务处理装饰器
    运行时间统计、异常处理、日志记录
    """
    if is_generator:

        def flow_decorator(
            handler: Callable[..., Generator[TaskResponse[Any]]],
        ):
            @wraps(handler)
            def wrapper(
                parameters: TaskBaseParameters, **kwargs
            ) -> Generator[TaskResponse[Any]]:
                logger.info(f"{task_name}({parameters['task_id']}) 开始执行")
                start_time = time.time()
                start_cpu_time = time.process_time()
                generator = handler(parameters, **kwargs)
                try:
                    for result in generator:
                        yield result
                    logger.info(f"{task_name}({parameters['task_id']}) 执行成功")
                except Exception as e:
                    logger.exception(
                        f"{task_name}({parameters['task_id']}) 执行失败: {e}"
                    )
                    logger.error(
                        f"{task_name}({parameters['task_id']}) 执行参数: {parameters}"
                    )
                    yield error_response(e)
                finally:
                    run_time = time.time() - start_time
                    cpu_time = time.process_time() - start_cpu_time
                    logger.info(
                        f"{task_name}({parameters['task_id']}) "
                        f"执行时间: {humanize.naturaldelta(run_time, minimum_unit='milliseconds')}({run_time * 1000}ms) "
                        f"(CPU: {humanize.naturaldelta(cpu_time, minimum_unit='milliseconds')})"
                    )
                    yield response("fee", {"run_time": run_time, "cpu_time": cpu_time})

            return wrapper

        return flow_decorator
    else:

        def decorator(handler: Callable[..., TaskResponse[Any]]):
            @wraps(handler)
            def wrapper(parameters: TaskBaseParameters, **kwargs) -> TaskResponse[Any]:
                start_time = time.time()
                start_cpu_time = time.process_time()
                try:
                    logger.info(f"{task_name}({parameters['task_id']}) 开始执行")
                    result = handler(parameters, **kwargs)
                    logger.info(f"{task_name}({parameters['task_id']}) 执行成功")
                except Exception as e:
                    logger.exception(
                        f"{task_name}({parameters['task_id']}) 执行失败: {e}"
                    )
                    logger.error(
                        f"{task_name}({parameters['task_id']}) 执行参数: {parameters}"
                    )
                    result = error_response(e)
                finally:
                    run_time = time.time() - start_time
                    cpu_time = time.process_time() - start_cpu_time
                    logger.info(
                        f"{task_name}({parameters['task_id']}) "
                        f"执行时间: {humanize.naturaldelta(run_time, minimum_unit='milliseconds')}({run_time * 1000}ms) "
                        f"(CPU: {humanize.naturaldelta(cpu_time, minimum_unit='milliseconds')})"
                    )
                    result.update(
                        {
                            "run_time": run_time,
                            "cpu_time": cpu_time,
                        }
                    )
                    return result

            return wrapper

        return decorator
