from abc import ABC, abstractmethod
from typing import List, Dict, Any, Generator, Union, Optional
import logging
import os
from pydantic import BaseModel, Field
from enum import Enum

__all__ = [
    'TaskStatus',
    'PROVIDER_STATUS_MAP',
    'map_provider_status',
    'is_terminal_status',
    'GenerationResult',
    'ServiceResponse',
    'BaseService',
    'AIProvider',
    'ImageProvider',
    'VideoProvider',
    'TextProvider',
    'AudioProvider',
    'MusicProvider',
    'BaseLLMService',
    'BaseImageService',
    'BaseVideoService',
    'BaseAudioService',
    'BaseMusicService',
]

logger = logging.getLogger(__name__)

class TaskStatus(str, Enum):
    """Unified task status enumeration for all services.
    
    This is the single source of truth for task statuses across the application.
    All provider-specific statuses should be mapped to these standard values.
    """ 
    PENDING = "pending"
    QUEUED = "queued"
    PROCESSING = "processing"
    SUCCEEDED = "succeeded"
    FAILED = "failed"
    TIMEOUT = "timeout"
    RETRYING = "retrying"
    CANCELLED = "cancelled"
    UNKNOWN = "unknown"


# Provider-specific status mappings
# Map provider status strings to unified TaskStatus
PROVIDER_STATUS_MAP: Dict[str, Dict[str, TaskStatus]] = {
    "dashscope": {
        "PENDING": TaskStatus.PENDING,
        "QUEUED": TaskStatus.QUEUED,
        "RUNNING": TaskStatus.PROCESSING,
        "SUCCEEDED": TaskStatus.SUCCEEDED,
        "FAILED": TaskStatus.FAILED,
        "UNKNOWN": TaskStatus.UNKNOWN,
    },
    "volcengine": {
        "queued": TaskStatus.QUEUED,
        "running": TaskStatus.PROCESSING,
        "succeeded": TaskStatus.SUCCEEDED,
        "failed": TaskStatus.FAILED,
        "cancelled": TaskStatus.CANCELLED,
        "expired": TaskStatus.FAILED,
    },
    "modelscope": {
        "PENDING": TaskStatus.PENDING,
        "RUNNING": TaskStatus.PROCESSING,
        "SUCCEED": TaskStatus.SUCCEEDED,  # 注意: ModelScope uses "SUCCEED" not "SUCCEEDED"
        "FAILED": TaskStatus.FAILED,
    },
}


def map_provider_status(provider: str, status: str) -> TaskStatus:
    """ 映射 provider-specific status string to unified TaskStatus.
    
    Args:
        provider: Provider name (e.g., 'dashscope', 'volcengine', 'modelscope')
        status: Provider-specific status string
        
    Returns:
        Unified TaskStatus enum value
        
    Example:
        >>> map_provider_status("dashscope", "SUCCEEDED")
        TaskStatus.SUCCEEDED
        >>> map_provider_status("volcengine", "running")
        TaskStatus.PROCESSING
    """ 
    provider_map = PROVIDER_STATUS_MAP.get(provider.lower(), {})
    mapped_status = provider_map.get(status)
    
    if mapped_status:
        return mapped_status
    
    # Fallback: try case-insensitive match against TaskStatus values
    status_lower = status.lower()
    for ts in TaskStatus:
        if ts.value == status_lower:
            return ts
    
    logger.warning(f"Unknown status '{status}' from provider '{provider}', defaulting to UNKNOWN")
    return TaskStatus.UNKNOWN


def is_terminal_status(status: TaskStatus) -> bool:
    """Check if a status is terminal (task will not change state anymore).
    
    Args:
        status: TaskStatus to check
        
    Returns:
        True if status is terminal, False otherwise
    """ 
    return status in (
        TaskStatus.SUCCEEDED,
        TaskStatus.FAILED,
        TaskStatus.TIMEOUT,
        TaskStatus.CANCELLED,
    )

class GenerationResult(BaseModel):
    """ 标准 result item for generation tasks."""
    url: Optional[str] = None
    content: Optional[str] = None
    orig_prompt: Optional[str] = None
    actual_prompt: Optional[str] = None
    usage: Optional[Dict[str, Any]] = None

class ServiceResponse(BaseModel):
    """ Standard response wrapper for AI services."""
    status: TaskStatus
    task_id: Optional[str] = None
    results: Optional[List[GenerationResult]] = None
    error: Optional[str] = None
    meta: Optional[Dict[str, Any]] = None

class AIProvider(ABC):
    """
    Abstract base class for all AI service providers.
    
    This is the top-level interface that all AI providers must implement.
    It defines the common contract for provider identification, health checks,
    and basic metadata.
    
    Requirement 7.1: Define abstract interfaces for each AI capability
    """
    
    @property
    @abstractmethod
    def provider_id(self) -> str:
        """Unique identifier for the provider (e.g., 'dashscope', 'kling', 'openai')"""
        pass
    
    @property
    @abstractmethod
    def provider_name(self) -> str:
        """Human-readable name for the provider"""
        pass
    
    @abstractmethod
    async def health_check(self) -> bool:
        """
        Check if the provider is healthy and available.
        
        Returns:
            True if healthy, False otherwise
        """
        pass
    
    @abstractmethod
    def get_capabilities(self) -> Dict[str, Any]:
        """Get provider capabilities and supported features.
        
        Returns:
            Dictionary of capabilities (e.g., {"supportsRefImage": True})
        """
        pass


class BaseService(AIProvider):
    """
    Base class for all AI services.

    Implements common functionality for all service types including
    health monitoring, configuration management, and model validation.

    Supports user-managed API keys with fallback to system keys:
    - User key (from UserApiKeyService) > System key (from config/env)

    Note: System keys are only available to superusers. Non-admin users must
    configure their own API keys in settings.
    """
    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        self.model_name = model_name
        self.api_key = api_key
        self.config: Dict[str, Any] = {}  # Will be injected by factory
        self._health_status = True
        self._last_health_check = None
        self._provider_id = kwargs.get('provider_id', 'unknown')
        self._provider_name = kwargs.get('provider_name', 'Unknown Provider')
        self._user_api_key_service = None  # Lazy loaded

    @property
    def provider_id(self) -> str:
        """Get provider ID from config or default"""
        if self.config:
            return self.config.get('provider', self._provider_id)
        return self._provider_id
    
    @property
    def provider_name(self) -> str:
        """Get provider name from config or default"""
        if self.config:
            return self.config.get('provider_name', self._provider_name)
        return self._provider_name

    async def health_check(self) -> bool:
        """
        Check if the service is healthy and available.
        Override this method in subclasses for specific health checks.
        """ 
        return self._health_status
    
    def get_capabilities(self) -> Dict[str, Any]:
        """Get service capabilities from config.
        Override in subclasses for dynamic capabilities.
        """
        if self.config:
            return self.config.get('capabilities', {})
        return {}
    
    def mark_unhealthy(self):
        """Mark service as unhealthy (e.g., after repeated failures)"""
        self._health_status = False
        logger.warning(f"Service {self.model_name} marked as unhealthy")
    
    def mark_healthy(self):
        """Mark service as healthy"""
        self._health_status = True
        logger.info(f"Service {self.model_name} marked as healthy")

    def set_user_api_key_service(self, service):
        """
        Set the user API key service for retrieving user-managed keys.

        Args:
            service: UserApiKeyService instance
        """
        self._user_api_key_service = service

    def _get_user_api_key_service(self):
        """
        Get or lazy-load the user API key service.

        Returns:
            UserApiKeyService instance or None
        """
        if self._user_api_key_service is None:
            try:
                from src.services.user_api_key_service import user_api_key_service
                self._user_api_key_service = user_api_key_service
            except ImportError:
                logger.debug("UserApiKeyService not available")
                return None
        return self._user_api_key_service

    # Provider ID → system env var name(s) for API key fallback
    _PROVIDER_ENV_MAP: Dict[str, str] = {
        'dashscope': 'DASHSCOPE_API_KEY',
        'volcengine': 'VOLCENGINE_API_KEY',
        'google': 'GOOGLE_API_KEY',
        'openai': 'OPENAI_API_KEY',
        'minimax': 'MINIMAX_API_KEY',
        'kling': 'KLING_ACCESS_KEY',
        'midjourney': 'MIDJOURNEY_API_KEY',
        'modelscope': 'MODELSCOPE_API_TOKEN',
    }

    def get_effective_api_key(self, user_id: Optional[str] = None) -> Optional[str]:
        """
        Get the effective API key for API calls.

        从系统环境变量获取 API Key。

        Returns:
            API key string or None
        """
        provider_id = self.provider_id

        env_var = self._PROVIDER_ENV_MAP.get(provider_id)
        if env_var:
            system_key = os.getenv(env_var)
            if system_key:
                return system_key

        logger.warning(f"No API key configured for provider {provider_id}")
        return None

    # 多密钥 provider 的系统 env var 映射（access_key + secret_key 等）
    _PROVIDER_EXTRA_ENV_MAP: Dict[str, Dict[str, str]] = {
        'kling': {'secret_key': 'KLING_SECRET_KEY'},
        'midjourney': {'api_secret': 'YOUCHUAN_SECRET_KEY', 'app_id': 'YOUCHUAN_APP_ID'},
    }

    def get_effective_api_key_with_config(
        self,
        user_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Get the effective API key with extra config for multi-key providers.

        从系统环境变量获取 API Key 及多密钥配置。
        适用于需要多个密钥的供应商（如 Kling: access+secret, Midjourney: app_id+secret）。

        Returns:
            Dict with keys like:
            - {"api_key": str} for single key
            - {"api_key": str, "extra_config": {"secret_key": ...}} for multi key
        """
        provider_id = self.provider_id

        env_var = self._PROVIDER_ENV_MAP.get(provider_id)
        if env_var:
            system_key = os.getenv(env_var)
            if system_key:
                result: Dict[str, Any] = {"api_key": system_key}
                # 补充多密钥 provider 的 secondary env vars
                extra_env = self._PROVIDER_EXTRA_ENV_MAP.get(provider_id)
                if extra_env:
                    extra_config = {}
                    for field, env_name in extra_env.items():
                        val = os.getenv(env_name)
                        if val:
                            extra_config[field] = val
                    if extra_config:
                        result["extra_config"] = extra_config
                return result

        logger.warning(f"No API key configured for provider {provider_id}")
        return {}

    def validate_model_from_config(self, _config_path: str, _service_category: str):
        """Validate if the model name is in the allowed list from a config file.
        Now uses ModelRegistry for validation instead of direct file scanning.
        """
        try:
            # Use ModelRegistry to check if model exists (works with new directory structure)
            from src.services.provider.registry import ModelRegistry
            
            config = ModelRegistry.get_config(self.model_name)
            if config:
                logger.info(f"Model {self.model_name} validated via ModelRegistry")
                return # Valid
            
            # If we reach here, model might be valid but not registered yet
            # (e.g., during initialization before registry is populated)
            logger.debug(f"Model {self.model_name} not found in registry (may be initializing)")

        except Exception as e:
            logger.warning(f"Failed to validate model name due to error: {e}")

class BaseLLMService(BaseService):
    """
    Base class for LLM services.
    
    Note: For standard OpenAI-compatible providers (OpenAI, Moonshot, DashScope, etc.),
    use OpenAIService from openai_service.py.
    """
    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key, **kwargs)

    @abstractmethod
    async def call(self, 
             messages: List[Dict[str, Any]], 
             stream: bool = True,
             temperature: float = 0.85,
             top_p: float = 0.8,
             **kwargs) -> ServiceResponse:
        pass


# ============================================================================
# Capability-Specific Provider Interfaces (Requirement 7.1)
# ============================================================================

class ImageProvider(ABC):
    """
    Abstract interface for image generation providers.
    
    All image generation services must implement this interface to ensure
    consistent behavior across different providers.
    
    Requirement 7.1: Define ImageProvider interface
    """
    
    @abstractmethod
    async def generate_image(
        self,
        prompt: str,
        negative_prompt: Optional[str] = None,
        size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        n: int = 1,
        **kwargs
    ) -> ServiceResponse:
        """Generate images from text prompt.
        
        Args:
            prompt: Text description of desired image
            negative_prompt: What to avoid in the image
            size: Image size (e.g., "1024*1024")
            aspect_ratio: Aspect ratio (e.g., "16:9")
            n: Number of images to generate
            **kwargs: Provider-specific parameters
            
        Returns:
            ServiceResponse with task_id or results
        """
        pass
    
    @abstractmethod
    async def generate_image_from_image(
        self,
        prompt: str,
        image_inputs: List[str],
        **kwargs
    ) -> ServiceResponse:
        """Generate images from reference images (I2I).
        
        Args:
            prompt: Text description
            image_inputs: List of reference image URLs or base64
            **kwargs: Provider-specific parameters
            
        Returns:
            ServiceResponse with task_id or results
        """
        pass


class VideoProvider(ABC):
    """
    Abstract interface for video generation providers.
    
    All video generation services must implement this interface to ensure
    consistent behavior across different providers.
    
    Requirement 7.1: Define VideoProvider interface
    """
    
    @abstractmethod
    async def generate_video_from_text(
        self,
        prompt: str,
        duration: int = 5,
        aspect_ratio: Optional[str] = None,
        **kwargs
    ) -> ServiceResponse:
        """Generate video from text prompt (T2V).
        
        Args:
            prompt: Text description of desired video
            duration: Video duration in seconds
            aspect_ratio: Aspect ratio (e.g., "16:9")
            **kwargs: Provider-specific parameters
            
        Returns:
            ServiceResponse with task_id
        """
        pass
    
    @abstractmethod
    async def generate_video_from_image(
        self,
        image: str,
        prompt: Optional[str] = None,
        duration: int = 5,
        **kwargs
    ) -> ServiceResponse:
        """Generate video from image (I2V).
        
        Args:
            image: Image URL or base64
            prompt: Optional text description
            duration: Video duration in seconds
            **kwargs: Provider-specific parameters
            
        Returns:
            ServiceResponse with task_id
        """
        pass


class TextProvider(ABC):
    """
    Abstract interface for text generation providers (LLM).
    
    All text generation services must implement this interface to ensure
    consistent behavior across different providers.
    
    Requirement 7.1: Define TextProvider interface
    """
    
    @abstractmethod
    async def generate_text(
        self,
        messages: List[Dict[str, Any]],
        temperature: float = 0.85,
        max_tokens: Optional[int] = None,
        stream: bool = False,
        **kwargs
    ) -> ServiceResponse:
        """Generate text from messages.
        
        Args:
            messages: List of message dicts with 'role' and 'content'
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            stream: Whether to stream the response
            **kwargs: Provider-specific parameters
            
        Returns:
            ServiceResponse with generated text
        """
        pass


class AudioProvider(ABC):
    """
    Abstract interface for audio generation providers.
    
    All audio generation services must implement this interface to ensure
    consistent behavior across different providers.
    
    Requirement 7.1: Define AudioProvider interface
    """
    
    @abstractmethod
    async def generate_audio(
        self,
        text: str,
        voice: str,
        format: str = "mp3",
        **kwargs
    ) -> ServiceResponse:
        """Generate audio from text (TTS).
        
        Args:
            text: Text to convert to speech
            voice: Voice identifier
            format: Audio format (mp3, wav, etc.)
            **kwargs: Provider-specific parameters
            
        Returns:
            ServiceResponse with audio URL or data
        """
        pass


class MusicProvider(ABC):
    """
    Abstract interface for music generation providers.

    All music generation services must implement this interface to ensure
    consistent behavior across different providers.
    """

    @abstractmethod
    async def generate_music(
        self,
        lyrics: str,
        prompt: Optional[str] = None,
        **kwargs
    ) -> ServiceResponse:
        """Generate music from lyrics and optional style prompt."""
        pass


# ============================================================================
# Concrete Base Service Classes
# ============================================================================

class BaseImageService(BaseService, ImageProvider):
    """
    Base class for image generation services.
    Implements both BaseService and ImageProvider interfaces.
    """
    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key, **kwargs)

    # ------------------------------------------------------------------
    # Frontend payload normalization helpers
    # ------------------------------------------------------------------
    def normalize_image_kwargs(self, **kwargs) -> Dict[str, Any]:
        """
        Normalize frontend payload into provider-friendly kwargs.

        Frontend currently passes a unified structure like:
            {
                "prompt": "...",
                "media": {
                    "images": [...],
                    "videos": [...],
                    "audio": "..."
                },
                ...
            }

        This helper:
        - Extracts media.images when present
        - Normalizes to image_inputs only
        """
        # Copy to avoid side effects on caller
        normalized: Dict[str, Any] = dict(kwargs)

        media = normalized.pop("media", None) or {}
        images = media.get("images") or normalized.get("image_inputs")

        if isinstance(images, str):
            images = [images]

        if images:
            normalized["image_inputs"] = images

        return normalized

    @abstractmethod
    async def generate(self, 
                       prompt: str, 
                       **kwargs) -> ServiceResponse:
        """Generate image content."""
        pass

    @abstractmethod
    async def check_status(self, task_id: str) -> ServiceResponse:
        """Check status of async generation task"""
        pass
    
    # Implement ImageProvider interface
    async def generate_image(
        self,
        prompt: str,
        negative_prompt: Optional[str] = None,
        size: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        n: int = 1,
        **kwargs
    ) -> ServiceResponse:
        """ 默认 implementation delegates to generate()""" 
        return await self.generate(
            prompt=prompt,
            negative_prompt=negative_prompt,
            size=size,
            aspect_ratio=aspect_ratio,
            n=n,
            **kwargs
        )
    
    async def generate_image_from_image(
        self,
        prompt: str,
        image_inputs: List[str],
        **kwargs
    ) -> ServiceResponse:
        """ 默认 implementation delegates to generate() with images""" 
        return await self.generate(
            prompt=prompt,
            image_inputs=image_inputs,
            **kwargs
        )


class BaseVideoService(BaseService, VideoProvider):
    """
    Base class for video generation services.
    Implements both BaseService and VideoProvider interfaces.
    """
    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key, **kwargs)

    @abstractmethod
    async def generate(self, 
                       prompt: str,
                       **kwargs) -> ServiceResponse:
        """Generate video content."""
        pass

    @abstractmethod
    async def check_status(self, task_id: str) -> ServiceResponse:
        """Check status of async generation task"""
        pass
    
    # Implement VideoProvider interface
    async def generate_video_from_text(
        self,
        prompt: str,
        duration: int = 5,
        aspect_ratio: Optional[str] = None,
        **kwargs
    ) -> ServiceResponse:
        """ 默认 implementation delegates to generate()""" 
        return await self.generate(
            prompt=prompt,
            duration=duration,
            aspect_ratio=aspect_ratio,
            **kwargs
        )
    
    async def generate_video_from_image(
        self,
        image: str,
        prompt: Optional[str] = None,
        duration: int = 5,
        **kwargs
    ) -> ServiceResponse:
        """ 默认 implementation delegates to generate() with image""" 
        return await self.generate(
            prompt=prompt or "",
            image_inputs=[image],
            duration=duration,
            **kwargs
        )


class BaseAudioService(BaseService, AudioProvider):
    """
    Base class for audio generation services.
    Implements both BaseService and AudioProvider interfaces.
    """
    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key, **kwargs)

    @abstractmethod
    async def generate(self, 
                       text: str, 
                       voice: str,
                       format: str,
                       **kwargs) -> ServiceResponse:
        """Generate audio content."""
        pass
    
    # Implement AudioProvider interface
    async def generate_audio(
        self,
        text: str,
        voice: str,
        format: str = "mp3",
        **kwargs
    ) -> ServiceResponse:
        """ 默认 implementation delegates to generate()""" 
        return await self.generate(
            text=text,
            voice=voice,
            format=format,
            **kwargs
        )


class BaseMusicService(BaseService, MusicProvider):
    """
    Base class for music generation services.
    Implements both BaseService and MusicProvider interfaces.
    """
    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key, **kwargs)

    @abstractmethod
    async def generate(
        self,
        lyrics: str,
        prompt: Optional[str] = None,
        **kwargs
    ) -> ServiceResponse:
        """Generate music content."""
        pass

    async def generate_music(
        self,
        lyrics: str,
        prompt: Optional[str] = None,
        **kwargs
    ) -> ServiceResponse:
        """Default implementation delegates to generate()."""
        return await self.generate(
            lyrics=lyrics,
            prompt=prompt,
            **kwargs
        )

    async def generate_lyrics(
        self,
        prompt: str,
        mode: str = "write_full_song",
        lyrics: Optional[str] = None,
        title: Optional[str] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Optional capability for providers that support dedicated lyrics generation.
        """
        raise NotImplementedError(
            f"{self.__class__.__name__} does not implement generate_lyrics()."
        )
