"""Tests for models/acestep/config.py"""

import json
import tempfile

from mlx_music.models.acestep.config import (
    AceStepConfig,
    AceStepDiTConfig,
    AceStepVAEConfig,
    load_config,
)


class TestAceStepDiTConfig:
    def test_defaults(self):
        cfg = AceStepDiTConfig()
        assert cfg.hidden_size == 2048
        assert cfg.num_hidden_layers == 24
        assert cfg.num_attention_heads == 16
        assert cfg.num_key_value_heads == 8
        assert cfg.head_dim == 128
        assert cfg.sliding_window == 128
        assert cfg.in_channels == 192
        assert cfg.patch_size == 2

    def test_layer_types_alternating(self):
        cfg = AceStepDiTConfig()
        assert len(cfg.layer_types) == 24
        assert cfg.layer_types[0] == "sliding_attention"
        assert cfg.layer_types[1] == "full_attention"
        assert cfg.layer_types[22] == "sliding_attention"
        assert cfg.layer_types[23] == "full_attention"


class TestAceStepVAEConfig:
    def test_defaults(self):
        cfg = AceStepVAEConfig()
        assert cfg.encoder_hidden_size == 128
        assert cfg.decoder_input_channels == 64
        assert cfg.audio_channels == 2
        assert cfg.sampling_rate == 48000

    def test_hop_length(self):
        cfg = AceStepVAEConfig()
        # Product of [2, 4, 4, 6, 10] = 1920
        assert cfg.hop_length == 1920

    def test_upsampling_ratios(self):
        cfg = AceStepVAEConfig()
        assert cfg.upsampling_ratios == [10, 6, 4, 4, 2]


class TestAceStepConfig:
    def test_latent_rate(self):
        cfg = AceStepConfig()
        # 48000 / 1920 = 25 Hz
        assert cfg.latent_rate == 25

    def test_duration_conversions(self):
        cfg = AceStepConfig()
        # 30 seconds = 750 frames at 25 Hz
        assert cfg.duration_to_latent_frames(30.0) == 750
        assert cfg.latent_frames_to_duration(750) == 30.0


class TestLoadConfig:
    def test_load_dit_config_json(self):
        """Load config from real HF format JSON."""
        dit_json = {
            "hidden_size": 2048,
            "num_hidden_layers": 24,
            "num_attention_heads": 16,
            "num_key_value_heads": 8,
            "head_dim": 128,
            "sliding_window": 128,
            "in_channels": 192,
            "is_turbo": True,
            "layer_types": ["sliding_attention", "full_attention"] * 12,
        }
        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
            json.dump(dit_json, f)
            path = f.name

        cfg = load_config(dit_config_path=path)
        assert cfg.dit.hidden_size == 2048
        assert cfg.dit.is_turbo is True
        assert len(cfg.dit.layer_types) == 24

    def test_load_vae_config_json(self):
        """Load VAE config from real HF format JSON."""
        vae_json = {
            "encoder_hidden_size": 128,
            "downsampling_ratios": [2, 4, 4, 6, 10],
            "decoder_input_channels": 64,
            "audio_channels": 2,
            "sampling_rate": 48000,
        }
        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
            json.dump(vae_json, f)
            path = f.name

        cfg = load_config(vae_config_path=path)
        assert cfg.vae.hop_length == 1920
        assert cfg.vae.audio_channels == 2
