"""
任务执行器模块

负责任务的实际执行、重试逻辑和错误处理。
"""

import asyncio
import logging
import time
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
from typing import Optional, Dict, Any, Callable, Awaitable

from prometheus_client import Counter, Histogram, Gauge

from src.services.provider.base import TaskStatus, ProviderService
from src.services.provider.registry import ModelRegistry

logger = logging.getLogger(__name__)


class ExecutionStatus(Enum):
    """执行状态"""
    SUCCESS = "success"
    FAILED = "failed"
    TIMEOUT = "timeout"
    RETRYABLE = "retryable"
    CANCELLED = "cancelled"


@dataclass
class TaskExecutionResult:
    """任务执行结果"""
    status: ExecutionStatus
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    execution_time: float = 0.0
    retry_count: int = 0
    provider_task_id: Optional[str] = None


@dataclass
class TaskConfig:
    """任务配置"""
    max_retries: int = 3
    retry_delay: int = 5  # seconds
    timeout: int = 300  # seconds
    exponential_backoff: bool = True
    max_concurrent: int = 5  # 最大并发数


# Prometheus 指标
EXECUTION_COUNTER = Counter(
    'task_execution_total',
    'Total task executions',
    ['task_type', 'status']
)

EXECUTION_DURATION = Histogram(
    'task_execution_duration_seconds',
    'Task execution duration',
    ['task_type', 'provider'],
    buckets=[1, 5, 10, 30, 60, 120, 300, 600]
)

ACTIVE_EXECUTIONS = Gauge(
    'task_active_executions',
    'Number of active task executions',
    ['task_type']
)

RETRY_COUNTER = Counter(
    'task_retry_total',
    'Total task retries',
    ['task_type', 'attempt']
)

PROVIDER_ERRORS = Counter(
    'task_provider_errors_total',
    'Provider API errors',
    ['provider', 'error_type']
)


class TaskExecutor:
    """
    任务执行器

    处理任务的实际执行，包括:
    - 调用提供商服务
    - 处理超时和重试
    - 错误处理和恢复
    - 并发控制
    """

    # 任务类型配置
    DEFAULT_CONFIGS: Dict[str, TaskConfig] = {
        "image": TaskConfig(max_retries=3, timeout=180, max_concurrent=10),
        "video": TaskConfig(max_retries=5, timeout=600, max_concurrent=5),
        "audio": TaskConfig(max_retries=3, timeout=240, max_concurrent=8),
        "music": TaskConfig(max_retries=3, timeout=300, max_concurrent=5),
        "script": TaskConfig(max_retries=2, timeout=120, max_concurrent=3),
    }

    def __init__(self):
        self._configs: Dict[str, TaskConfig] = dict(self.DEFAULT_CONFIGS)
        self._semaphores: Dict[str, asyncio.Semaphore] = {}
        self._running_tasks: Dict[str, asyncio.Task] = {}
        self._model_registry = ModelRegistry()

        # 初始化信号量
        for task_type, config in self._configs.items():
            self._semaphores[task_type] = asyncio.Semaphore(config.max_concurrent)

        logger.info("TaskExecutor initialized")

    def update_config(self, task_type: str, config: TaskConfig):
        """更新任务类型配置"""
        self._configs[task_type] = config
        self._semaphores[task_type] = asyncio.Semaphore(config.max_concurrent)

    async def execute(
        self,
        task_id: str,
        task_type: str,
        provider: str,
        model: str,
        params: Dict[str, Any],
        retry_count: int = 0,
        cancellation_event: Optional[asyncio.Event] = None
    ) -> TaskExecutionResult:
        """
        执行任务

        Args:
            task_id: 任务 ID
            task_type: 任务类型
            provider: 提供商名称
            model: 模型标识
            params: 任务参数
            retry_count: 当前重试次数
            cancellation_event: 取消事件

        Returns:
            TaskExecutionResult
        """
        config = self._configs.get(task_type, TaskConfig())
        semaphore = self._semaphores.get(task_type)

        start_time = time.time()
        ACTIVE_EXECUTIONS.labels(task_type=task_type).inc()

        try:
            async with semaphore:
                # 检查取消
                if cancellation_event and cancellation_event.is_set():
                    return TaskExecutionResult(
                        status=ExecutionStatus.CANCELLED,
                        execution_time=time.time() - start_time
                    )

                # 获取提供商服务
                service = ModelRegistry.get(provider)
                if not service:
                    return TaskExecutionResult(
                        status=ExecutionStatus.FAILED,
                        error=f"Provider service not found: {provider}",
                        execution_time=time.time() - start_time
                    )

                # 注入用户 API key 服务
                from src.services.user_api_key_service import user_api_key_service
                service.set_user_api_key_service(user_api_key_service)

                # 执行任务
                result = await self._execute_with_timeout(
                    service=service,
                    task_type=task_type,
                    model=model,
                    params=params,
                    timeout=config.timeout,
                    cancellation_event=cancellation_event
                )

                execution_time = time.time() - start_time

                # 更新指标
                EXECUTION_DURATION.labels(
                    task_type=task_type,
                    provider=provider
                ).observe(execution_time)

                EXECUTION_COUNTER.labels(
                    task_type=task_type,
                    status=result.status.value
                ).inc()

                result.execution_time = execution_time
                return result

        except asyncio.TimeoutError:
            execution_time = time.time() - start_time
            EXECUTION_COUNTER.labels(
                task_type=task_type,
                status="timeout"
            ).inc()

            return TaskExecutionResult(
                status=ExecutionStatus.TIMEOUT,
                error=f"Task timed out after {config.timeout}s",
                execution_time=execution_time
            )

        except Exception as e:
            execution_time = time.time() - start_time
            logger.exception(f"Task execution error: {e}")

            EXECUTION_COUNTER.labels(
                task_type=task_type,
                status="error"
            ).inc()

            return TaskExecutionResult(
                status=ExecutionStatus.FAILED,
                error=str(e),
                execution_time=execution_time
            )

        finally:
            ACTIVE_EXECUTIONS.labels(task_type=task_type).dec()
            self._running_tasks.pop(task_id, None)

    async def _execute_with_timeout(
        self,
        service,
        task_type: str,
        model: str,
        params: Dict[str, Any],
        timeout: int,
        cancellation_event: Optional[asyncio.Event]
    ) -> TaskExecutionResult:
        """
        带超时的任务执行

        Args:
            service: 提供商服务实例
            task_type: 任务类型
            model: 模型标识
            params: 任务参数
            timeout: 超时时间
            cancellation_event: 取消事件

        Returns:
            TaskExecutionResult
        """
        try:
            # 从 params 中提取 user_id 并调用相应的服务方法
            # user_id 用于获取用户管理的 API key
            user_id = params.get("user_id")

            if task_type == "image":
                # 判断是文生图还是图生图
                if params.get("image_inputs") or params.get("images"):
                    result = await asyncio.wait_for(
                        service.generate_image_from_image(
                            prompt=params.get("prompt", ""),
                            image_inputs=params.get("image_inputs") or params.get("images", []),
                            **{k: v for k, v in params.items() if k not in ["prompt", "image_inputs", "images", "user_id"]}
                        ),
                        timeout=timeout
                    )
                else:
                    result = await asyncio.wait_for(
                        service.generate_image(
                            prompt=params.get("prompt", ""),
                            **{k: v for k, v in params.items() if k not in ["prompt", "user_id"]}
                        ),
                        timeout=timeout
                    )
            elif task_type == "video":
                result = await asyncio.wait_for(
                    service.generate(
                        prompt=params.get("prompt", ""),
                        **{k: v for k, v in params.items() if k not in ["prompt", "user_id"]}
                    ),
                    timeout=timeout
                )
            elif task_type == "audio":
                result = await asyncio.wait_for(
                    service.generate(
                        text=params.get("text", ""),
                        voice=params.get("voice", "alloy"),
                        format=params.get("format", "mp3"),
                        **{k: v for k, v in params.items() if k not in ["text", "voice", "format", "user_id"]}
                    ),
                    timeout=timeout
                )
            elif task_type == "music":
                result = await asyncio.wait_for(
                    service.generate(
                        lyrics=params.get("lyrics", ""),
                        prompt=params.get("prompt"),
                        **{k: v for k, v in params.items() if k not in ["lyrics", "prompt", "user_id"]}
                    ),
                    timeout=timeout
                )
            elif task_type == "script":
                result = await asyncio.wait_for(
                    service.call(
                        messages=params.get("messages", []),
                        **{k: v for k, v in params.items() if k not in ["messages", "user_id"]}
                    ),
                    timeout=timeout
                )
            else:
                return TaskExecutionResult(
                    status=ExecutionStatus.FAILED,
                    error=f"Unknown task type: {task_type}"
                )

            # 检查取消
            if cancellation_event and cancellation_event.is_set():
                return TaskExecutionResult(
                    status=ExecutionStatus.CANCELLED
                )

            # Convert ServiceResponse to dict
            result_dict = result.model_dump() if hasattr(result, 'model_dump') else result

            return TaskExecutionResult(
                status=ExecutionStatus.SUCCESS,
                result=result_dict,
                provider_task_id=result_dict.get("task_id") if isinstance(result_dict, dict) else None
            )

        except asyncio.TimeoutError:
            raise

        except Exception as e:
            error_msg = str(e)

            # 判断是否可重试
            is_retryable = self._is_retryable_error(e)

            if is_retryable:
                return TaskExecutionResult(
                    status=ExecutionStatus.RETRYABLE,
                    error=error_msg
                )

            return TaskExecutionResult(
                status=ExecutionStatus.FAILED,
                error=error_msg
            )

    def _is_retryable_error(self, error: Exception) -> bool:
        """
        判断错误是否可重试

        Args:
            error: 异常对象

        Returns:
            是否可重试
        """
        error_str = str(error).lower()

        # 可重试的错误类型
        retryable_patterns = [
            "rate limit",
            "too many requests",
            "timeout",
            "connection",
            "network",
            "temporary",
            "unavailable",
            "503",
            "429",
            "502",
            "504",
        ]

        return any(pattern in error_str for pattern in retryable_patterns)

    async def should_retry(
        self,
        task_type: str,
        retry_count: int,
        result: TaskExecutionResult
    ) -> bool:
        """
        判断是否应该重试

        Args:
            task_type: 任务类型
            retry_count: 当前重试次数
            result: 执行结果

        Returns:
            是否应该重试
        """
        config = self._configs.get(task_type, TaskConfig())

        # 检查重试次数
        if retry_count >= config.max_retries:
            return False

        # 检查状态
        if result.status == ExecutionStatus.RETRYABLE:
            RETRY_COUNTER.labels(task_type=task_type, attempt=retry_count + 1).inc()
            return True

        return False

    def calculate_retry_delay(self, task_type: str, retry_count: int) -> float:
        """
        计算重试延迟

        Args:
            task_type: 任务类型
            retry_count: 当前重试次数

        Returns:
            延迟时间（秒）
        """
        config = self._configs.get(task_type, TaskConfig())

        if config.exponential_backoff:
            # 指数退避: delay * 2^retry_count
            delay = config.retry_delay * (2 ** retry_count)
            # 添加 jitter 防止惊群
            import random
            delay += random.uniform(0, 1)
            return min(delay, 300)  # 最大 5 分钟

        return config.retry_delay

    async def cancel_task(self, task_id: str) -> bool:
        """
        取消正在执行的任务

        Args:
            task_id: 任务 ID

        Returns:
            是否成功取消
        """
        task = self._running_tasks.get(task_id)
        if task and not task.done():
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
            return True
        return False

    def get_stats(self) -> Dict[str, Any]:
        """获取执行器统计信息"""
        return {
            "active_tasks": len(self._running_tasks),
            "configs": {
                task_type: {
                    "max_retries": config.max_retries,
                    "timeout": config.timeout,
                    "max_concurrent": config.max_concurrent
                }
                for task_type, config in self._configs.items()
            }
        }
