import os
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 _get_rss_mb() -> float:
    """获取当前进程实时 RSS（MB）"""
    try:
        import psutil
        return psutil.Process().memory_info().rss / (1024 * 1024)
    except Exception:
        return 0.0


def _extract_task_profile(parameters: dict) -> str:
    """提取影响内存的关键任务参数摘要"""
    try:
        sc = parameters.get("strategy_config", {})
        dr = sc.get("date_range", {})
        hm = sc.get("hold_method", {})
        hrv = hm.get("hold_range_value", {})
        bt = sc.get("backtest_type", "?")
        parts = [
            f"date={dr.get('start_date', '?')}~{dr.get('end_date', '?')}",
            f"hold={hrv.get('start', '?')}~{hrv.get('end', '?')}",
            f"bt_type={bt}",
            f"factors={len(sc.get('rank_factors', []))}r/{len(sc.get('filter_factors', []))}f/{len(sc.get('computed_factors', []))}c",
        ]
        pools = sc.get("pools", [])
        if pools:
            parts.append(f"pools={len(pools)}")
        return " ".join(parts)
    except Exception:
        return "parse_error"


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]]:
                rss_before = _get_rss_mb()
                profile = _extract_task_profile(parameters)
                logger.info(
                    f"{task_name}({parameters['task_id']}) 开始执行 "
                    f"[RSS={rss_before:.0f}MB] [{profile}]"
                )
                start_time = time.time()
                start_cpu_time = time.process_time()
                generator = handler(parameters, **kwargs)
                try:
                    for result in generator:
                        yield result
                    rss_after = _get_rss_mb()
                    logger.info(
                        f"{task_name}({parameters['task_id']}) 执行成功 "
                        f"[RSS={rss_after:.0f}MB, delta=+{rss_after - rss_before:.0f}MB]"
                    )
                except Exception as e:
                    rss_err = _get_rss_mb()
                    logger.exception(
                        f"{task_name}({parameters['task_id']}) 执行失败: {e} "
                        f"[RSS={rss_err:.0f}MB, delta=+{rss_err - rss_before:.0f}MB]"
                    )
                    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]:
                rss_before = _get_rss_mb()
                profile = _extract_task_profile(parameters)
                start_time = time.time()
                start_cpu_time = time.process_time()
                try:
                    logger.info(
                        f"{task_name}({parameters['task_id']}) 开始执行 "
                        f"[RSS={rss_before:.0f}MB] [{profile}]"
                    )
                    result = handler(parameters, **kwargs)
                    rss_after = _get_rss_mb()
                    logger.info(
                        f"{task_name}({parameters['task_id']}) 执行成功 "
                        f"[RSS={rss_after:.0f}MB, delta=+{rss_after - rss_before:.0f}MB]"
                    )
                except Exception as e:
                    rss_err = _get_rss_mb()
                    logger.exception(
                        f"{task_name}({parameters['task_id']}) 执行失败: {e} "
                        f"[RSS={rss_err:.0f}MB, delta=+{rss_err - rss_before:.0f}MB]"
                    )
                    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
