# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from typing import Any, List

from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig


@dataclass
class LTXVideoVAEArchConfig(VAEArchConfig):
    # Architecture params
    in_channels: int = 3
    latent_channels: int = 128
    out_channels: int = 3
    temporal_compression_ratio: int = 8
    spatial_compression_ratio: int = 32
    block_out_channels: List[int] = field(
        default_factory=lambda: [256, 512, 1024, 2048]
    )
    down_block_types: List[str] = field(
        default_factory=lambda: [
            "LTX2VideoDownBlock3D",
            "LTX2VideoDownBlock3D",
            "LTX2VideoDownBlock3D",
            "LTX2VideoDownBlock3D",
        ]
    )
    spatio_temporal_scaling: List[bool] = field(
        default_factory=lambda: [True, True, True, True]
    )
    layers_per_block: List[int] = field(default_factory=lambda: [4, 6, 6, 2, 2])
    downsample_type: List[str] = field(
        default_factory=lambda: [
            "spatial",
            "temporal",
            "spatiotemporal",
            "spatiotemporal",
        ]
    )
    patch_size: int = 4
    patch_size_t: int = 1
    resnet_norm_eps: float = 1e-6
    encoder_causal: bool = True
    encoder_spatial_padding_mode: str = "zeros"

    decoder_block_out_channels: List[int] = field(
        default_factory=lambda: [256, 512, 1024]
    )
    decoder_spatio_temporal_scaling: List[bool] = field(
        default_factory=lambda: [True, True, True]
    )
    decoder_layers_per_block: List[int] = field(default_factory=lambda: [5, 5, 5, 5])
    decoder_causal: bool = False
    decoder_spatial_padding_mode: str = "reflect"

    # Native LTX variant metadata.
    ltx_variant: str = "ltx_2"
    condition_encoder_subdir: str = ""
    video_encoder_variant: str = "ltx_2"
    video_encoder_config: dict[str, Any] = field(default_factory=dict)
    video_decoder_variant: str = "ltx_2"
    video_decoder_config: dict[str, Any] = field(default_factory=dict)


@dataclass
class LTXVideoVAEConfig(VAEConfig):
    arch_config: LTXVideoVAEArchConfig = field(default_factory=LTXVideoVAEArchConfig)
    auto_parallel_decode_min_latent_elements_per_rank: int = 1024

    def auto_parallel_decode_prefers_spatial_shard(self) -> bool:
        return True
