"""
Pydantic models for API requests and responses.
Inspired by OpenAI's API but tailored for Qwen3-TTS.
"""

from typing import List, Optional, Dict, Any
from enum import Enum
from pydantic import BaseModel, Field


class TTSModel(str, Enum):
    """Available TTS models."""
    BASE = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
    # Add more model variants as needed


class VoiceCloneMode(str, Enum):
    """Voice cloning modes."""
    DISABLED = "disabled"
    REFERENCE_AUDIO = "reference_audio"
    CUSTOM = "custom"


class StreamOptions(BaseModel):
    """Streaming configuration options."""
    emit_every_frames: int = Field(default=8, ge=1, le=100, description="Emit chunks every N frames")
    decode_window_frames: int = Field(default=80, ge=1, le=256, description="Decode window size in frames")
    overlap_samples: int = Field(default=512, ge=0, le=2048, description="Overlap samples for smooth transitions")


class TTSRequest(BaseModel):
    """Base TTS request model."""
    text: str = Field(..., min_length=1, max_length=10000, description="Text to synthesize")
    model: TTSModel = Field(default=TTSModel.BASE, description="Model to use")
    language: str = Field(default="English", description="Language for TTS")
    voice_clone_mode: VoiceCloneMode = Field(default=VoiceCloneMode.DISABLED, description="Voice cloning mode")
    speed: float = Field(default=1.0, ge=0.5, le=2.0, description="Speech speed multiplier")
    pitch: float = Field(default=1.0, ge=0.5, le=2.0, description="Pitch adjustment")


class StreamingTTSRequest(TTSRequest):
    """Request model for streaming TTS."""
    stream_options: StreamOptions = Field(default_factory=StreamOptions, description="Streaming configuration")


class TTSResponse(BaseModel):
    """Response model for TTS requests."""
    id: str = Field(..., description="Request ID")
    object: str = Field(default="audio", description="Object type")
    created: int = Field(..., description="Unix timestamp of creation")
    model: str = Field(..., description="Model used")
    audio_base64: str = Field(..., description="Audio data in base64 format")
    duration: float = Field(..., description="Duration in seconds")
    sample_rate: int = Field(default=24000, description="Sample rate in Hz")
    language: str = Field(..., description="Language used")


class ErrorResponse(BaseModel):
    """Error response model."""
    error: str = Field(..., description="Error message")
    code: str = Field(..., description="Error code")
    details: Optional[Dict[str, Any]] = Field(default=None, description="Additional error details")


class HealthResponse(BaseModel):
    """Health check response."""
    status: str = Field(default="healthy", description="Health status")
    model: str = Field(..., description="Current model")
    device: str = Field(..., description="Computation device")
    ready: bool = Field(default=True, description="Ready for requests")


class ModelInfoResponse(BaseModel):
    """Model information response."""
    id: str = Field(..., description="Model ID")
    object: str = Field(default="model", description="Object type")
    owned_by: str = Field(default="Alibaba", description="Model owner")
    supported_languages: List[str] = Field(..., description="Supported languages")
    supports_streaming: bool = Field(default=True, description="Supports streaming")
    supports_voice_clone: bool = Field(default=True, description="Supports voice cloning")


class ModelsListResponse(BaseModel):
    """List of available models."""
    object: str = Field(default="list", description="Object type")
    data: List[ModelInfoResponse] = Field(..., description="List of models")


# ============== Voice Management ==============

class VoiceResponse(BaseModel):
    """Custom voice information."""
    id: str = Field(..., description="Voice ID")
    name: str = Field(..., description="Voice name")
    language: str = Field(..., description="Voice language")
    created_at: int = Field(..., description="Creation timestamp")
    object: str = Field(default="voice", description="Object type")


class VoicesListResponse(BaseModel):
    """List of custom voices."""
    object: str = Field(default="list", description="Object type")
    data: List[VoiceResponse] = Field(..., description="List of voices")


# ============== Text Validation ==============

class TextValidationRequest(BaseModel):
    """Request to validate text for TTS."""
    text: str = Field(..., min_length=1, max_length=10000, description="Text to validate")
    language: str = Field(default="English", description="Language for validation")


class TextValidationResponse(BaseModel):
    """Text validation response."""
    valid: bool = Field(..., description="Whether text is valid for TTS")
    language: str = Field(..., description="Detected language")
    character_count: int = Field(..., description="Number of characters")
    estimated_duration: float = Field(..., description="Estimated audio duration in seconds")
    warnings: List[str] = Field(default_factory=list, description="Any warnings about the text")


# ============== Batch Processing ==============

class BatchJobStatus(str, Enum):
    """Batch job status."""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class BatchItem(BaseModel):
    """Single item in a batch request."""
    text: str = Field(..., min_length=1, description="Text to synthesize")
    language: str = Field(default="English", description="Language")
    voice_clone_mode: VoiceCloneMode = Field(default=VoiceCloneMode.DISABLED, description="Voice cloning mode")


class CreateBatchRequest(BaseModel):
    """Request to create a batch job."""
    items: List[BatchItem] = Field(..., min_items=1, max_items=100, description="Batch items")
    model: TTSModel = Field(default=TTSModel.BASE, description="Model to use")


class BatchJobResponse(BaseModel):
    """Batch job information."""
    id: str = Field(..., description="Job ID")
    object: str = Field(default="batch", description="Object type")
    status: BatchJobStatus = Field(..., description="Current status")
    created_at: int = Field(..., description="Creation timestamp")
    updated_at: int = Field(..., description="Last update timestamp")
    request_counts: Dict[str, int] = Field(..., description="Counts: total, processing, completed, failed")
    output_file_id: Optional[str] = Field(default=None, description="Output file ID when completed")


class BatchResultItem(BaseModel):
    """Single result item from batch job."""
    index: int = Field(..., description="Item index")
    status: str = Field(..., description="Item status: success or error")
    error: Optional[str] = Field(default=None, description="Error message if failed")
    audio_base64: Optional[str] = Field(default=None, description="Generated audio in base64")
    duration: Optional[float] = Field(default=None, description="Audio duration")


class BatchResultsResponse(BaseModel):
    """Batch job results."""
    job_id: str = Field(..., description="Batch job ID")
    object: str = Field(default="batch.results", description="Object type")
    status: BatchJobStatus = Field(..., description="Job status")
    data: List[BatchResultItem] = Field(..., description="Result items")


# ============== Usage & Quota ==============

class UsageResponse(BaseModel):
    """API usage statistics."""
    object: str = Field(default="usage", description="Object type")
    requests_made: int = Field(..., description="Total requests made")
    audio_generated_seconds: float = Field(..., description="Total audio seconds generated")
    audio_generated_minutes: float = Field(..., description="Total audio minutes generated")
    requests_by_model: Dict[str, int] = Field(..., description="Requests per model")
    requests_by_language: Dict[str, int] = Field(..., description="Requests per language")


class QuotaResponse(BaseModel):
    """Rate limits and quotas."""
    object: str = Field(default="quota", description="Object type")
    requests_per_minute: int = Field(..., description="Rate limit (requests/minute)")
    max_text_length: int = Field(..., description="Maximum text length")
    max_batch_size: int = Field(..., description="Maximum batch items")
    concurrent_requests: int = Field(..., description="Max concurrent requests")
    remaining_requests: int = Field(..., description="Remaining requests in current window")


# ============== Audio Conversion ==============

class AudioFormat(str, Enum):
    """Supported audio formats."""
    WAV = "wav"
    MP3 = "mp3"
    OGG = "ogg"
    FLAC = "flac"


class ConvertAudioRequest(BaseModel):
    """Request to convert audio format/sample rate."""
    audio_base64: str = Field(..., description="Audio data in base64")
    target_format: AudioFormat = Field(default=AudioFormat.WAV, description="Target audio format")
    target_sample_rate: Optional[int] = Field(default=None, ge=8000, le=48000, description="Target sample rate (Hz)")


class ConvertAudioResponse(BaseModel):
    """Audio conversion response."""
    id: str = Field(..., description="Request ID")
    object: str = Field(default="audio", description="Object type")
    created: int = Field(..., description="Timestamp")
    format: str = Field(..., description="Output format")
    sample_rate: int = Field(..., description="Output sample rate")
    audio_base64: str = Field(..., description="Converted audio in base64")
    duration: float = Field(..., description="Duration in seconds")


# ============== Model Configuration ==============

class ModelConfigResponse(BaseModel):
    """Model configuration details."""
    id: str = Field(..., description="Model ID")
    object: str = Field(default="model.config", description="Object type")
    size: str = Field(..., description="Model size (parameters)")
    tokenizer_type: str = Field(..., description="Tokenizer type")
    languages: List[str] = Field(..., description="Supported languages")
    max_text_length: int = Field(..., description="Maximum input text length")
    output_sample_rate: int = Field(..., description="Output sample rate (Hz)")
    streaming_supported: bool = Field(default=True, description="Streaming support")
    voice_cloning_supported: bool = Field(default=True, description="Voice cloning support")
    recommended_parameters: Dict[str, Any] = Field(..., description="Recommended generation parameters")