"""
Generation-related schemas.
"""
from typing import List, Optional, Dict, Any, Literal
from datetime import datetime
from pydantic import BaseModel, Field, field_validator, model_validator, ConfigDict
import uuid
import re

from .common import GenerationStatus
from src.utils.validators import VALID_IMAGE_MODELS


class BaseGenerationRecord(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: float = Field(default_factory=lambda: datetime.now().timestamp())
    status: GenerationStatus
    model: str
    prompt: str
    result_url: Optional[str] = Field(None, alias="resultUrl")
    result_urls: Optional[List[str]] = Field(None, alias="resultUrls")
    error: Optional[str] = None

    model_config = ConfigDict(populate_by_name=True)


class ImageGenerationParams(BaseModel):
    negative_prompt: Optional[str] = Field(None, alias="negativePrompt")
    aspect_ratio: Optional[str] = Field(None, alias="aspectRatio")
    width: Optional[int] = None
    height: Optional[int] = None
    dimension: Optional[str] = None
    image_inputs: Optional[List[str]] = Field(None, alias="imageInputs")
    loras: Optional[List[Dict[str, Any]]] = None
    model_params: Optional[Dict[str, Any]] = Field(None, alias="modelConfig")

    model_config = ConfigDict(populate_by_name=True, protected_namespaces=())


class ImageGenerationRecord(BaseGenerationRecord):
    type: Literal['image'] = 'image'
    params: ImageGenerationParams


class VideoGenerationParams(BaseModel):
    negative_prompt: Optional[str] = Field(None, alias="negativePrompt")
    duration: Optional[float] = None
    aspect_ratio: Optional[str] = Field(None, alias="aspectRatio")
    resolution: Optional[str] = None

    # Unified Media Inputs
    image_inputs: Optional[List[str]] = Field(None, alias="imageInputs")
    video_inputs: Optional[List[str]] = Field(None, alias="videoInputs")
    audio_inputs: Optional[List[str]] = Field(None, alias="audioInputs")

    enable_audio: Optional[bool] = Field(None, alias="enableAudio")
    audio_url: Optional[str] = Field(None, alias="audioUrl")

    camera_control: Optional[str] = Field(None, alias="cameraControl")
    shot_type: Optional[str] = Field(None, alias="shotType")
    model_params: Optional[Dict[str, Any]] = Field(None, alias="modelConfig")

    model_config = ConfigDict(populate_by_name=True, protected_namespaces=())


class VideoGenerationRecord(BaseGenerationRecord):
    type: Literal['video'] = 'video'
    params: VideoGenerationParams


GenerationRecord = ImageGenerationRecord | VideoGenerationRecord


class StandardVideoParams(BaseModel):
    """Standard video generation parameters for adapter system."""
    prompt: Optional[str] = None
    negative_prompt: Optional[str] = Field(None, alias="negativePrompt")

    # Input media
    image_urls: Optional[List[str]] = Field(None, alias="imageUrls", description="Input images for I2V")
    video_urls: Optional[List[str]] = Field(None, alias="videoUrls", description="Input videos for V2V")
    audio_urls: Optional[List[str]] = Field(None, alias="audioUrls", description="Input audio for Audio2Video")

    first_frame_url: Optional[str] = Field(None, alias="firstFrameUrl", description="First frame image")
    last_frame_url: Optional[str] = Field(None, alias="lastFrameUrl", description="Last frame image")
    reference_video_urls: Optional[List[str]] = Field(None, alias="referenceVideoUrls", description="Reference videos for R2V")

    # Output controls
    duration: Optional[int] = Field(None, description="Video duration in seconds")
    ratio: Optional[str] = Field(None, description="Aspect ratio e.g. 16:9, 9:16")
    resolution: Optional[str] = Field(None, description="Resolution level e.g. 720P, 1080P")
    size: Optional[str] = Field(None, description="Resolution string e.g. 1280*720")
    fps: Optional[int] = Field(None, description="Frames per second")

    # Generation controls
    guidance_scale: Optional[float] = Field(None, alias="guidanceScale")

    # Audio
    generate_audio: Optional[bool] = Field(None, alias="generateAudio")
    audio_url: Optional[str] = Field(None, alias="audioUrl")

    # Camera motion
    camera_motion: Optional[Dict[str, Any]] = Field(None, alias="cameraMotion")

    # Advanced
    loras: Optional[List[Dict[str, Any]]] = None
    extra_params: Optional[Dict[str, Any]] = Field(None, alias="extraParams")

    model_config = ConfigDict(populate_by_name=True)


class StandardImageParams(BaseModel):
    """Standard image generation parameters for adapter system."""
    prompt: str
    negative_prompt: Optional[str] = Field(None, alias="negativePrompt")

    # Input media
    reference_image_urls: Optional[List[str]] = Field(None, alias="referenceImageUrls", description="Reference images for I2I")

    # Output controls
    size: Optional[str] = Field(None, description="Size string e.g. 1024*1024")
    width: Optional[int] = None
    height: Optional[int] = None
    aspect_ratio: Optional[str] = Field(None, alias="aspectRatio")
    n: int = Field(1, description="Number of images to generate")

    # Generation controls
    steps: Optional[int] = Field(None, description="Number of inference steps")
    guidance_scale: Optional[float] = Field(None, alias="guidanceScale")

    # Advanced
    loras: Optional[List[Dict[str, Any]]] = None
    extra_params: Optional[Dict[str, Any]] = Field(None, alias="extraParams")

    model_config = ConfigDict(populate_by_name=True)


class ImageGenerationRequest(BaseModel):
    """Generic Image Generation Request"""
    model_config = ConfigDict(populate_by_name=True)

    prompt: str
    negative_prompt: Optional[str] = Field(None, alias="negativePrompt")
    model: str  # Composite ID format: provider/model_key

    # Unified naming: Media Inputs
    image_inputs: Optional[List[str]] = Field(None, alias="imageInputs")

    # Unified naming: Dimensions
    resolution: Optional[str] = Field(None, description="Resolution level e.g. 1K, 2K, 4K")
    aspect_ratio: Optional[str] = Field(None, alias="aspectRatio", description="Aspect ratio e.g. 16:9")

    # Generation count
    n: int = 1

    # Optional Context
    project_id: Optional[str] = Field(None, alias="projectId")
    source: Optional[str] = None
    source_id: Optional[str] = Field(None, alias="sourceId")

    # Other experimental parameters
    extra_params: Optional[Dict[str, Any]] = Field(None, alias="extraParams")

    @field_validator('model')
    @classmethod
    def validate_model_format(cls, v: str) -> str:
        """Validate model format must be provider/model_key"""
        if '/' not in v and v in VALID_IMAGE_MODELS:
            return v

        if '/' not in v:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model must be in format 'provider/model_key', got: '{v}'. "
                f"Example: 'dashscope/qwen-image'"
            )

        parts = v.split('/')
        if len(parts) != 2:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model format invalid: '{v}'. Must have exactly one '/' separator."
            )

        provider, model_key = parts
        if not provider or not model_key:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model format invalid: '{v}'. Both provider and model_key must be non-empty."
            )

        return v

    @field_validator('prompt')
    @classmethod
    def validate_prompt(cls, v: str) -> str:
        """Validate prompt is not empty"""
        if not v or not v.strip():
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException("prompt", "Prompt cannot be empty")
        return v.strip()

    @field_validator('n')
    @classmethod
    def validate_n(cls, v: int) -> int:
        """Validate generation count"""
        if v < 1 or v > 10:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException("n", "Number of images must be between 1 and 10")
        return v

    @field_validator('aspect_ratio')
    @classmethod
    def validate_aspect_ratio(cls, v: Optional[str]) -> Optional[str]:
        """Validate aspect ratio format"""
        if v is None:
            return v
        if not re.match(r'^\d+:\d+$', v):
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "aspect_ratio",
                f"Invalid aspect ratio format: {v}. Expected format: 'width:height' (e.g., '16:9')"
            )
        return v

    @field_validator('resolution')
    @classmethod
    def validate_resolution(cls, v: Optional[str]) -> Optional[str]:
        """Validate resolution format"""
        if v is None:
            return v

        quality_pattern = r'^(1K|2K|4K|720P|1080P|480P|360P)$'

        if not re.match(quality_pattern, v, re.IGNORECASE):
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "resolution",
                f"Invalid resolution format: {v}. Expected quality level (e.g., '1K', '2K', '720P')"
            )
        return v


class VideoGenerationRequest(BaseModel):
    """Generic Video Generation Request"""
    model_config = ConfigDict(populate_by_name=True)

    prompt: str
    negative_prompt: Optional[str] = Field(None, alias="negativePrompt", description="Negative prompt")
    model: str  # Composite ID format: provider/model_key

    # Unified naming: Media Inputs
    image_inputs: Optional[List[str]] = Field(None, alias="imageInputs", description="Input image list")
    video_inputs: Optional[List[str]] = Field(None, alias="videoInputs", description="Input video list")
    audio_inputs: Optional[List[str]] = Field(None, alias="audioInputs", description="Input audio list")

    # Unified naming: Controls
    enable_audio: Optional[bool] = Field(True, alias="enableAudio", description="Generate audio")
    duration: int = 5

    # Unified naming: Dimensions/Quality
    aspect_ratio: Optional[str] = Field(None, alias="aspectRatio")
    resolution: Optional[str] = Field(None, description="Quality level e.g. 720P, 1080P")

    # Generation count
    n: int = Field(1, description="Number of videos to generate")

    # Advanced Controls
    camera_control: Optional[str] = Field(None, alias="cameraControl", description="Camera control type")
    shot_type: Optional[str] = Field(None, alias="shotType", description="Shot type")

    # Context
    project_id: Optional[str] = Field(None, alias="projectId")
    source: Optional[str] = None
    source_id: Optional[str] = Field(None, alias="sourceId")

    # Other experimental parameters
    extra_params: Optional[Dict[str, Any]] = Field(None, alias="extraParams")

    @field_validator('model')
    @classmethod
    def validate_model_format(cls, v: str) -> str:
        """Validate model format must be provider/model_key"""
        if '/' not in v:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model must be in format 'provider/model_key', got: '{v}'. "
                f"Example: 'dashscope/wan2.6-video'"
            )

        parts = v.split('/')
        if len(parts) != 2:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model format invalid: '{v}'. Must have exactly one '/' separator."
            )

        provider, model_key = parts
        if not provider or not model_key:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model format invalid: '{v}'. Both provider and model_key must be non-empty."
            )

        return v

    @field_validator('duration')
    @classmethod
    def validate_duration(cls, v: int) -> int:
        """Validate video duration"""
        if v < 1 or v > 60:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException("duration", "Duration must be between 1 and 60 seconds")
        return v

    @field_validator('aspect_ratio')
    @classmethod
    def validate_aspect_ratio(cls, v: Optional[str]) -> Optional[str]:
        """Validate aspect ratio format"""
        if v is None:
            return v
        if not re.match(r'^\d+:\d+$', v):
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "aspect_ratio",
                f"Invalid aspect ratio format: {v}. Expected format: 'width:height' (e.g., '16:9')"
            )
        return v

    @field_validator('n')
    @classmethod
    def validate_n(cls, v: int) -> int:
        """Validate generation count"""
        if v < 1 or v > 10:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException("n", "Number of videos must be between 1 and 10")
        return v

    @model_validator(mode='after')
    def validate_input_required(self) -> 'VideoGenerationRequest':
        """Validate at least one input is provided"""
        if not self.prompt and not self.image_inputs and not self.video_inputs:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "prompt/imageInputs/videoInputs",
                "Either prompt, image_inputs, or video_inputs must be provided"
            )
        return self


class AudioGenerationRequest(BaseModel):
    """Generic Audio Generation Request"""
    model_config = ConfigDict(populate_by_name=True)

    text: str
    model: str  # Composite ID format: provider/model_key

    # Audio-specific parameters
    voice: Optional[str] = None

    # Context
    project_id: Optional[str] = Field(None, alias="projectId")
    storyboard_id: Optional[str] = Field(None, alias="storyboardId")
    source: Optional[str] = None
    source_id: Optional[str] = Field(None, alias="sourceId")

    # Other experimental parameters
    extra_params: Optional[Dict[str, Any]] = Field(None, alias="extraParams")

    @field_validator('model')
    @classmethod
    def validate_model_format(cls, v: str) -> str:
        """Validate model format must be provider/model_key"""
        if '/' not in v:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model must be in format 'provider/model_key', got: '{v}'. "
                f"Example: 'volcengine/doubao-tts'"
            )

        parts = v.split('/')
        if len(parts) != 2:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model format invalid: '{v}'. Must have exactly one '/' separator."
            )

        provider, model_key = parts
        if not provider or not model_key:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model format invalid: '{v}'. Both provider and model_key must be non-empty."
            )

        return v

    @field_validator('text')
    @classmethod
    def validate_text(cls, v: str) -> str:
        """Validate text is not empty"""
        if not v or not v.strip():
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException("text", "Text cannot be empty")
        return v.strip()


class MusicGenerationRequest(BaseModel):
    """Generic Music Generation Request"""
    model_config = ConfigDict(populate_by_name=True)

    model: str  # Composite ID format: provider/model_key
    generation_mode: Literal["music", "lyrics"] = Field("music", alias="mode", description="Unified interface mode")
    prompt: Optional[str] = Field(None, description="Music style/mood/scene description")
    lyrics: Optional[str] = Field(None, description="Lyrics content")
    lyrics_prompt: Optional[str] = Field(None, alias="lyricsPrompt", description="Theme for auto-generated lyrics")
    title: Optional[str] = Field(None, description="Title for lyrics generation")
    lyrics_mode: Optional[Literal["write_full_song", "edit"]] = Field("write_full_song", alias="lyricsMode")
    seed_lyrics: Optional[str] = Field(None, alias="seedLyrics", description="Original lyrics for edit mode")

    # Music output controls
    output_format: Optional[str] = Field("url", alias="outputFormat", description="Output format: url or hex")
    stream: Optional[bool] = False
    audio_setting: Optional[Dict[str, Any]] = Field(None, alias="audioSetting")
    aigc_watermark: Optional[bool] = Field(False, alias="aigcWatermark")

    # Context
    project_id: Optional[str] = Field(None, alias="projectId")
    source: Optional[str] = None
    source_id: Optional[str] = Field(None, alias="sourceId")

    # Other experimental parameters
    extra_params: Optional[Dict[str, Any]] = Field(None, alias="extraParams")

    @field_validator('model')
    @classmethod
    def validate_model_format(cls, v: str) -> str:
        """Validate model format must be provider/model_key"""
        if '/' not in v:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model must be in format 'provider/model_key', got: '{v}'. "
                f"Example: 'minimax/music-2.5'"
            )

        parts = v.split('/')
        if len(parts) != 2:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model format invalid: '{v}'. Must have exactly one '/' separator."
            )

        provider, model_key = parts
        if not provider or not model_key:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "model",
                f"Model format invalid: '{v}'. Both provider and model_key must be non-empty."
            )

        return v

    @field_validator('output_format')
    @classmethod
    def validate_output_format(cls, v: Optional[str]) -> Optional[str]:
        if v is None:
            return v
        val = v.lower()
        if val not in {"url", "hex"}:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException("output_format", "output_format must be 'url' or 'hex'")
        return val

    @model_validator(mode='after')
    def validate_lyrics_or_prompt(self) -> 'MusicGenerationRequest':
        """Validate unified interface parameters based on mode"""
        has_lyrics = bool(self.lyrics and self.lyrics.strip())
        has_lyrics_prompt = bool(self.lyrics_prompt and self.lyrics_prompt.strip())
        has_prompt = bool(self.prompt and self.prompt.strip())

        if self.generation_mode == "lyrics":
            if not has_prompt:
                from src.utils.errors import InvalidParameterException
                raise InvalidParameterException("prompt", "prompt is required when mode='lyrics'")
            if self.lyrics_mode == "edit" and not (self.seed_lyrics and self.seed_lyrics.strip()):
                from src.utils.errors import InvalidParameterException
                raise InvalidParameterException("seedLyrics", "seedLyrics is required when lyricsMode='edit'")
            return self

        # mode=music
        if not has_lyrics and not has_lyrics_prompt and not has_prompt:
            from src.utils.errors import InvalidParameterException
            raise InvalidParameterException(
                "lyrics/lyricsPrompt/prompt",
                "Either lyrics or lyricsPrompt (or prompt for auto lyrics) must be provided for mode='music'"
            )
        return self


class GenerationResponse(BaseModel):
    task_id: Optional[str]
    status: str
    raw_response: Optional[str] = None


class UpscaleImageRequest(BaseModel):
    image_url: str = Field(..., description="Original image URL")
    rate: int = Field(2, description="Upscale rate (default: 2)")
    extra_params: Optional[Dict[str, Any]] = Field(None, alias="extraParams")


class UpscaleImageResponse(BaseModel):
    original_url: str
    upscaled_url: str


class GenerateAudioRequest(BaseModel):
    text: str
    voice: Optional[str] = None
    model: Optional[str] = None
    provider: Optional[str] = None
    project_id: Optional[str] = None
    storyboard_id: Optional[str] = None
    extra_params: Optional[Dict[str, Any]] = Field(None, alias="extraParams")


class AudioResponse(BaseModel):
    id: str
    audio_url: str
    status: str


class ExportProjectRequest(BaseModel):
    project_id: str
    format: str = "mp4"


class ExportResponse(BaseModel):
    download_url: str
