"""ACE-Step 1.5 model configuration for MLX."""

from dataclasses import dataclass, field
from typing import List, Optional
import json
from pathlib import Path


@dataclass
class AceStepDiTConfig:
    """Configuration for the ACE-Step DiT (Diffusion Transformer)."""

    hidden_size: int = 2048
    intermediate_size: int = 6144
    num_hidden_layers: int = 24
    num_attention_heads: int = 16
    num_key_value_heads: int = 8
    head_dim: int = 128
    hidden_act: str = "silu"
    max_position_embeddings: int = 32768
    rms_norm_eps: float = 1e-6
    rope_theta: float = 1000000.0
    attention_bias: bool = False
    use_sliding_window: bool = True
    sliding_window: int = 128
    vocab_size: int = 64003

    # Layer types: alternating sliding/full attention
    layer_types: List[str] = field(
        default_factory=lambda: [
            "sliding_attention" if i % 2 == 0 else "full_attention" for i in range(24)
        ]
    )

    # Input/output dimensions
    in_channels: int = 192  # 64 (acoustic) + 64 (src_latent) + 64 (chunk_mask)
    audio_acoustic_hidden_dim: int = 64
    patch_size: int = 2

    # Text/lyric/timbre encoder config
    text_hidden_dim: int = 1024
    num_lyric_encoder_hidden_layers: int = 8
    timbre_hidden_dim: int = 64
    num_timbre_encoder_hidden_layers: int = 4
    timbre_fix_frame: int = 750

    # Pooler/tokenizer config
    num_attention_pooler_hidden_layers: int = 2
    num_audio_decoder_hidden_layers: int = 24
    pool_window_size: int = 5

    # FSQ config
    fsq_dim: int = 2048
    fsq_input_levels: List[int] = field(default_factory=lambda: [8, 8, 8, 5, 5, 5])
    fsq_input_num_quantizers: int = 1

    # Timestep sampling
    timestep_mu: float = -0.4
    timestep_sigma: float = 1.0
    data_proportion: float = 0.5

    # Model variant
    model_version: str = "turbo"
    is_turbo: bool = True


@dataclass
class AceStepVAEConfig:
    """Configuration for the AutoencoderOobleck (audio VAE)."""

    encoder_hidden_size: int = 128
    downsampling_ratios: List[int] = field(default_factory=lambda: [2, 4, 4, 6, 10])
    channel_multiples: List[int] = field(default_factory=lambda: [1, 2, 4, 8, 16])
    decoder_channels: int = 128
    decoder_input_channels: int = 64  # latent dim (splits to 2x32 for mean/scale)
    audio_channels: int = 2  # stereo
    sampling_rate: int = 48000

    @property
    def hop_length(self) -> int:
        """Total downsampling factor = product of all ratios."""
        result = 1
        for r in self.downsampling_ratios:
            result *= r
        return result

    @property
    def upsampling_ratios(self) -> List[int]:
        """Decoder upsampling = reversed encoder downsampling."""
        return self.downsampling_ratios[::-1]

    @property
    def latent_dim(self) -> int:
        """Latent dimension = decoder_input_channels."""
        return self.decoder_input_channels


@dataclass
class AceStepConfig:
    """Top-level configuration for the complete ACE-Step model."""

    dit: AceStepDiTConfig = field(default_factory=AceStepDiTConfig)
    vae: AceStepVAEConfig = field(default_factory=AceStepVAEConfig)

    # Text encoder model (loaded via mlx-lm)
    text_encoder_model: str = "Qwen/Qwen3-Embedding-0.6B"

    # LM planner models (loaded via mlx-lm)
    lm_planner_models: List[str] = field(
        default_factory=lambda: [
            "ACE-Step/acestep-5Hz-lm-0.6B",
            "ACE-Step/acestep-5Hz-lm-1.7B",
            "ACE-Step/acestep-5Hz-lm-4B",
        ]
    )

    # Audio settings
    sample_rate: int = 48000
    latent_rate: int = 25  # 48000 / 1920 = 25 Hz

    @property
    def latent_dim(self) -> int:
        return self.vae.decoder_input_channels

    def samples_to_latent_frames(self, num_samples: int) -> int:
        """Convert audio sample count to latent frame count."""
        return num_samples // self.vae.hop_length

    def duration_to_latent_frames(self, duration_seconds: float) -> int:
        """Convert duration in seconds to latent frame count."""
        return int(duration_seconds * self.latent_rate)

    def latent_frames_to_duration(self, num_frames: int) -> float:
        """Convert latent frame count to duration in seconds."""
        return num_frames / self.latent_rate


def load_config(
    dit_config_path: Optional[str] = None, vae_config_path: Optional[str] = None
) -> AceStepConfig:
    """Load config from JSON files (HuggingFace format).

    Args:
        dit_config_path: Path to DiT config.json
        vae_config_path: Path to VAE config.json

    Returns:
        AceStepConfig with values from files, defaults for missing fields.
    """
    dit_cfg = AceStepDiTConfig()
    vae_cfg = AceStepVAEConfig()

    if dit_config_path is not None:
        with open(dit_config_path) as f:
            d = json.load(f)
        # Map HF config keys to our dataclass fields
        for key in [
            "hidden_size",
            "intermediate_size",
            "num_hidden_layers",
            "num_attention_heads",
            "num_key_value_heads",
            "head_dim",
            "hidden_act",
            "max_position_embeddings",
            "rms_norm_eps",
            "rope_theta",
            "attention_bias",
            "use_sliding_window",
            "sliding_window",
            "layer_types",
            "in_channels",
            "audio_acoustic_hidden_dim",
            "patch_size",
            "text_hidden_dim",
            "num_lyric_encoder_hidden_layers",
            "timbre_hidden_dim",
            "num_timbre_encoder_hidden_layers",
            "timbre_fix_frame",
            "num_attention_pooler_hidden_layers",
            "num_audio_decoder_hidden_layers",
            "pool_window_size",
            "fsq_dim",
            "fsq_input_levels",
            "fsq_input_num_quantizers",
            "timestep_mu",
            "timestep_sigma",
            "data_proportion",
            "model_version",
            "vocab_size",
        ]:
            if key in d:
                setattr(dit_cfg, key, d[key])
        if "is_turbo" in d:
            dit_cfg.is_turbo = d["is_turbo"]

    if vae_config_path is not None:
        with open(vae_config_path) as f:
            d = json.load(f)
        for key in [
            "encoder_hidden_size",
            "downsampling_ratios",
            "channel_multiples",
            "decoder_channels",
            "decoder_input_channels",
            "audio_channels",
            "sampling_rate",
        ]:
            if key in d:
                setattr(vae_cfg, key, d[key])

    return AceStepConfig(dit=dit_cfg, vae=vae_cfg)
