""" 轮询 Strategies for Async Task Management

This module provides configurable polling strategies for checking
the status of asynchronous AI generation tasks. Different strategies
can be used based on the expected response time characteristics of
different providers and task types.

Usage:
    config = PollingConfig(strategy="exponential_backoff", base_interval=2)
    strategy = create_polling_strategy(config)
    
    for attempt in range(config.max_attempts):
        interval = strategy.get_next_interval(attempt)
        await asyncio.sleep(interval)
        # Check task status...
"""

from abc import ABC, abstractmethod
from typing import Optional, Dict, Any, Type
from dataclasses import dataclass, field
from enum import Enum
import logging

logger = logging.getLogger(__name__)


class PollingStrategyType(str, Enum):
    """Available polling strategy types."""
    FIXED = "fixed"
    LINEAR = "linear"
    EXPONENTIAL_BACKOFF = "exponential_backoff"
    ADAPTIVE = "adaptive"


@dataclass
class PollingConfig:
    """ 配置 for polling behavior.
    
    This can be specified in model configuration JSON files:
    
    {
        "polling_config": {
            "strategy": "exponential_backoff",
            "max_attempts": 90,
            "base_interval": 2,
            "max_interval": 30,
            "timeout": 300
        }
    }
    """
    strategy: str = "exponential_backoff"
    max_attempts: int = 60
    base_interval: float = 2.0  # seconds
    max_interval: float = 30.0  # seconds
    timeout: int = 300  # seconds (5 minutes)
    
    # Linear strategy
    increment: float = 1.0  # seconds to add each attempt
    
    # Exponential backoff strategy
    multiplier: float = 2.0
    
    # Early stop conditions
    early_stop_statuses: list = field(default_factory=lambda: ["succeeded", "failed", "cancelled", "timeout"])
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "PollingConfig":
        """ 创建 PollingConfig from dictionary.
        
        Args:
            data: Configuration dictionary
            
        Returns:
            PollingConfig instance
        """
        return cls(
            strategy=data.get("strategy", "exponential_backoff"),
            max_attempts=data.get("max_attempts", 60),
            base_interval=data.get("base_interval", 2.0),
            max_interval=data.get("max_interval", 30.0),
            timeout=data.get("timeout", 300),
            increment=data.get("increment", 1.0),
            multiplier=data.get("multiplier", 2.0),
            early_stop_statuses=data.get("early_stop_statuses", ["succeeded", "failed", "cancelled", "timeout"])
        )
    
    def to_dict(self) -> Dict[str, Any]:
        """ 转换 to dictionary for serialization.
        
        Returns:
            Dictionary representation of polling configuration
        """
        return {
            "strategy": self.strategy,
            "max_attempts": self.max_attempts,
            "base_interval": self.base_interval,
            "max_interval": self.max_interval,
            "timeout": self.timeout,
            "increment": self.increment,
            "multiplier": self.multiplier,
            "early_stop_statuses": self.early_stop_statuses
        }


class PollingStrategy(ABC):
    """Base class for polling strategies."""
    
    def __init__(self, config: PollingConfig):
        self.config = config
    
    @abstractmethod
    def get_next_interval(self, attempt: int) -> float:
        """Calculate the interval before the next poll.
        
        Args:
            attempt: Current attempt number (0-indexed)
            
        Returns:
            Interval in seconds before next poll
        """
        pass
    
    def should_continue(self, attempt: int, elapsed_time: float) -> bool:
        """Check if polling should continue.
        
        Args:
            attempt: Current attempt number
            elapsed_time: Total elapsed time in seconds
            
        Returns:
            True if should continue polling, False otherwise
        """
        if attempt >= self.config.max_attempts:
            logger.warning(f"Max polling attempts ({self.config.max_attempts}) reached")
            return False
        
        if elapsed_time >= self.config.timeout:
            logger.warning(f"Polling timeout ({self.config.timeout}s) reached")
            return False
        
        return True


class FixedIntervalStrategy(PollingStrategy):
    """Poll at fixed intervals.
    
    Good for:
    - Tasks with predictable completion times
    - Simple implementation requirements
    """
    
    def get_next_interval(self, attempt: int) -> float:
        return self.config.base_interval


class LinearStrategy(PollingStrategy):
    """Increase polling interval linearly.
    
    Interval = base_interval + (attempt * increment)
    
    Good for:
    - Tasks that may complete quickly but sometimes take longer
    - Moderate reduction in API calls
    """
    
    def get_next_interval(self, attempt: int) -> float:
        interval = self.config.base_interval + (attempt * self.config.increment)
        return min(interval, self.config.max_interval)


class ExponentialBackoffStrategy(PollingStrategy):
    """Exponentially increase polling interval.
    
    Interval = base_interval * (multiplier ^ attempt)
    
    Good for:
    - Tasks with highly variable completion times
    - Reducing load on provider APIs
    - Long-running tasks
    """
    
    def get_next_interval(self, attempt: int) -> float:
        interval = self.config.base_interval * (self.config.multiplier ** attempt)
        return min(interval, self.config.max_interval)


class AdaptiveStrategy(PollingStrategy):
    """Adaptive polling based on observed patterns.
    
    Starts with short intervals and adapts based on
    whether the task shows progress.
    
    Good for:
    - Unknown task characteristics
    - Production environments with diverse workloads
    """
    
    def __init__(self, config: PollingConfig):
        super().__init__(config)
        self._consecutive_no_change = 0
    
    def get_next_interval(self, attempt: int) -> float:
        # 启动 aggressive, then back off
        if attempt < 5:
            return self.config.base_interval
        elif attempt < 15:
            return self.config.base_interval * 1.5
        elif attempt < 30:
            return self.config.base_interval * 2
        else:
            return min(self.config.base_interval * 3, self.config.max_interval)
    
    def record_progress(self, has_progress: bool):
        """ 记录 whether the last check showed progress.
        
        Args:
            has_progress: Whether the task showed any progress
        """
        if has_progress:
            self._consecutive_no_change = 0
        else:
            self._consecutive_no_change += 1


# Strategy Registry
STRATEGY_REGISTRY: Dict[str, Type[PollingStrategy]] = {
    PollingStrategyType.FIXED.value: FixedIntervalStrategy,
    PollingStrategyType.LINEAR.value: LinearStrategy,
    PollingStrategyType.EXPONENTIAL_BACKOFF.value: ExponentialBackoffStrategy,
    PollingStrategyType.ADAPTIVE.value: AdaptiveStrategy,
}


def create_polling_strategy(config: Optional[PollingConfig] = None) -> PollingStrategy:
    """ 创建 a polling strategy from configuration.
    
    Args:
        config: Polling configuration. If None, uses default exponential backoff.
        
    Returns:
        PollingStrategy instance
        
    Raises:
        ValueError: If strategy type is unknown
    """
    if config is None:
        config = PollingConfig()
    
    strategy_class = STRATEGY_REGISTRY.get(config.strategy.lower())
    
    if strategy_class is None:
        logger.warning(f"Unknown strategy '{config.strategy}', falling back to exponential_backoff")
        strategy_class = ExponentialBackoffStrategy
    
    return strategy_class(config)


# 默认 configurations for different task types
DEFAULT_POLLING_CONFIGS: Dict[str, PollingConfig] = {
    "image": PollingConfig(
        strategy="exponential_backoff",
        max_attempts=30,
        base_interval=2.0,
        max_interval=10.0,
        timeout=180,  # 3 minutes
    ),
    "video": PollingConfig(
        strategy="exponential_backoff",
        max_attempts=90,
        base_interval=3.0,
        max_interval=30.0,
        timeout=600,  # 10 minutes
    ),
    "audio": PollingConfig(
        strategy="fixed",
        max_attempts=20,
        base_interval=2.0,
        timeout=60,  # 1 minute
    ),
}


def get_default_polling_config(task_type: str) -> PollingConfig:
    """ 获取 default polling configuration for a task type.
    
    Args:
        task_type: Type of task ('image', 'video', 'audio')
        
    Returns:
        Default PollingConfig for that task type
    """
    return DEFAULT_POLLING_CONFIGS.get(task_type, PollingConfig())
