# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Cosmos3 sound tokenizer integration."""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import torch
from vllm.logger import init_logger

from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.progress_bar import _is_rank_zero

from .audio_tokenizer import Cosmos3AVAEAudioTokenizer

logger = init_logger(__name__)

DEFAULT_SOUND_SAMPLE_RATE = 48000
DEFAULT_SOUND_CHANNELS = 2
DEFAULT_SOUND_DIM = 64
DEFAULT_SOUND_HOP_SIZE = 1920
DEFAULT_SOUND_LATENT_FPS = DEFAULT_SOUND_SAMPLE_RATE / DEFAULT_SOUND_HOP_SIZE
DEFAULT_SOUND_NORMALIZE_LATENTS = False
DEFAULT_SOUND_NORMALIZATION_TYPE = "none"
DEFAULT_SOUND_TANH_INPUT_SCALE = 1.5
DEFAULT_SOUND_TANH_OUTPUT_SCALE = 3.5
DEFAULT_SOUND_TANH_CLAMP = 0.995
SOUND_TOKENIZER_COMPONENT_NAME = "sound_tokenizer"
SOUND_TOKENIZER_CHECKPOINT_NAME = "diffusion_pytorch_model.safetensors"


def _pipeline_args(od_config: OmniDiffusionConfig) -> dict[str, Any]:
    return dict(getattr(od_config, "custom_pipeline_args", None) or {})


def _config_get(config: Any, key: str, default: Any = None) -> Any:
    if config is None:
        return default
    if isinstance(config, dict):
        return config.get(key, default)
    if hasattr(config, "get"):
        value = config.get(key, None)
        return default if value is None else value
    return getattr(config, key, default)


def _config_path_get(config: Any, *keys: str) -> Any:
    value = config
    for key in keys:
        value = _config_get(value, key, None)
        if value is None:
            return None
    return value


def _sound_tokenizer_config_from(config: Any) -> Any:
    """Return nested ``sound_tokenizer`` config from Cosmos3 config shapes."""
    for path in (
        ("sound_tokenizer",),
        ("model", "config", "sound_tokenizer"),
        ("config", "sound_tokenizer"),
        ("model_config", "sound_tokenizer"),
    ):
        value = _config_path_get(config, *path)
        if value is not None:
            return value
    return None


def _nested_sound_tokenizer_configs(od_config: OmniDiffusionConfig | None) -> tuple[Any, ...]:
    if od_config is None:
        return ()
    configs = []
    for source in (
        getattr(od_config, "model_config", None),
        getattr(od_config, "tf_model_config", None),
    ):
        config = _sound_tokenizer_config_from(source)
        if config is not None:
            configs.append(config)
    return tuple(configs)


def _first_value_from_configs(configs: tuple[Any, ...], keys: tuple[str, ...]) -> Any:
    for config in configs:
        for key in keys:
            value = _config_get(config, key, None)
            if value is not None:
                return value
    return None


def _top_level_model_value(od_config: OmniDiffusionConfig | None, keys: tuple[str, ...]) -> Any:
    if od_config is None:
        return None
    for source in (
        getattr(od_config, "model_config", None),
        getattr(od_config, "tf_model_config", None),
    ):
        for key in keys:
            for path in ((key,), ("model", "config", key), ("config", key), ("model_config", key)):
                value = _config_path_get(source, *path)
                if value is not None:
                    return value
    return None


def _custom_arg_value(args: dict[str, Any], keys: tuple[str, ...]) -> Any:
    for key in keys:
        value = args.get(key)
        if value is not None:
            return value
    return None


def _as_bool(value: Any) -> bool:
    if isinstance(value, str):
        return value.strip().lower() in {"1", "true", "yes", "on"}
    return bool(value)


def _as_audio_channels(value: Any) -> int:
    if isinstance(value, bool):
        return 2 if value else 1
    if isinstance(value, str) and value.strip().lower() in {
        "1",
        "0",
        "true",
        "false",
        "yes",
        "no",
        "on",
        "off",
    }:
        return 2 if _as_bool(value) else 1
    return int(value)


def _resolve_model_file(path: Any, model_root: str | None) -> str | None:
    if not path:
        return None
    path = str(path)
    if "://" in path or os.path.isabs(path) or os.path.exists(path) or not model_root:
        return path
    return str(Path(model_root) / path)


def _load_sound_tokenizer_component_config(config_path: str | None) -> dict[str, Any]:
    if not config_path:
        return {}
    with open(config_path, encoding="utf-8") as f:
        config = json.load(f)
    if not isinstance(config, dict):
        raise TypeError(f"Cosmos3 sound tokenizer config must be a JSON object, got {type(config)!r}.")
    return config


def _component_audio_channels(config: dict[str, Any]) -> Any:
    if config.get("dec_out_channels") is not None:
        return config["dec_out_channels"]
    if config.get("audio_channels") is not None:
        return config["audio_channels"]
    if config.get("stereo") is not None:
        return 2 if _as_bool(config["stereo"]) else 1
    return None


def _component_arch_values(config: dict[str, Any]) -> dict[str, Any]:
    values = {
        "sample_rate": config.get("sampling_rate", config.get("sample_rate")),
        "audio_channels": _component_audio_channels(config),
        "io_channels": config.get("vocoder_input_dim", config.get("io_channels", config.get("latent_ch"))),
        "hop_size": config.get("hop_size"),
    }
    return {key: value for key, value in values.items() if value is not None}


def _resolve_arch_value(
    od_config: OmniDiffusionConfig,
    args: dict[str, Any],
    component_values: dict[str, Any],
    *,
    field: str,
    custom_keys: tuple[str, ...],
    nested_keys: tuple[str, ...],
    top_level_keys: tuple[str, ...],
    default: Any,
    cast,
) -> Any:
    custom_value = _custom_arg_value(args, custom_keys)
    component_value = component_values.get(field)
    if component_value is not None:
        resolved = cast(component_value)
        if custom_value is not None and cast(custom_value) != resolved:
            raise ValueError(
                "Conflicting Cosmos3 sound tokenizer architecture override for "
                f"{field}: component config has {resolved!r}, custom args have {cast(custom_value)!r}."
            )
        return resolved

    if custom_value is not None:
        return cast(custom_value)

    nested_value = _first_value_from_configs(_nested_sound_tokenizer_configs(od_config), nested_keys)
    if nested_value is not None:
        return cast(nested_value)

    top_value = _top_level_model_value(od_config, top_level_keys)
    if top_value is not None:
        return cast(top_value)

    return cast(default)


def _resolve_normalization_value(
    od_config: OmniDiffusionConfig,
    args: dict[str, Any],
    *,
    name: str,
    default: Any,
    aliases: tuple[str, ...] = (),
) -> Any:
    keys = (f"sound_{name}", name, *aliases)
    custom_value = _custom_arg_value(args, keys)
    if custom_value is not None:
        return custom_value
    nested_value = _first_value_from_configs(_nested_sound_tokenizer_configs(od_config), (name, *aliases))
    return default if nested_value is None else nested_value


def get_sound_config_value(
    od_config: OmniDiffusionConfig,
    name: str,
    default: Any,
    aliases: tuple[str, ...] = (),
) -> Any:
    # Backward-compatible generic accessor.  Prefer the more specific helpers
    # below for Cosmos3 sound tokenizer fields so precedence stays explicit.
    keys = (name, *aliases)
    for config in (
        _pipeline_args(od_config),
        getattr(od_config, "model_config", None),
        getattr(od_config, "tf_model_config", None),
    ):
        if config is None:
            continue
        for key in keys:
            if hasattr(config, "get"):
                value = config.get(key, None)
            else:
                value = getattr(config, key, None)
            if value is not None:
                return value
    return default


def get_sound_sample_rate(od_config: OmniDiffusionConfig) -> int:
    args = _pipeline_args(od_config)
    return _resolve_arch_value(
        od_config,
        args,
        {},
        field="sample_rate",
        custom_keys=("sound_sample_rate", "sample_rate"),
        nested_keys=("sample_rate", "sampling_rate"),
        top_level_keys=("sound_sample_rate", "sample_rate"),
        default=DEFAULT_SOUND_SAMPLE_RATE,
        cast=int,
    )


def get_sound_channels(od_config: OmniDiffusionConfig) -> int:
    args = _pipeline_args(od_config)
    return _resolve_arch_value(
        od_config,
        args,
        {},
        field="audio_channels",
        custom_keys=("sound_audio_channels", "audio_channels", "stereo"),
        nested_keys=("audio_channels", "dec_out_channels", "stereo"),
        top_level_keys=("sound_audio_channels", "audio_channels", "stereo"),
        default=DEFAULT_SOUND_CHANNELS,
        cast=_as_audio_channels,
    )


def get_sound_hop_size(od_config: OmniDiffusionConfig) -> int:
    args = _pipeline_args(od_config)
    return _resolve_arch_value(
        od_config,
        args,
        {},
        field="hop_size",
        custom_keys=("sound_hop_size", "hop_size"),
        nested_keys=("hop_size",),
        top_level_keys=("sound_hop_size", "hop_size"),
        default=DEFAULT_SOUND_HOP_SIZE,
        cast=int,
    )


class Cosmos3SoundTokenizer:
    """Thin adapter around the local AVAE tokenizer implementation."""

    def __init__(self, tokenizer: Any) -> None:
        self.tokenizer = tokenizer
        self.sample_rate = int(getattr(tokenizer, "sample_rate", DEFAULT_SOUND_SAMPLE_RATE))
        self.audio_channels = int(getattr(tokenizer, "audio_channels", DEFAULT_SOUND_CHANNELS))
        self.latent_ch = int(getattr(tokenizer, "latent_ch", DEFAULT_SOUND_DIM))
        self.hop_size = int(getattr(tokenizer, "temporal_compression_factor", DEFAULT_SOUND_HOP_SIZE))
        if self.hop_size <= 0:
            raise ValueError(f"Cosmos3 sound tokenizer hop_size must be positive, got {self.hop_size}.")
        self.latent_fps = float(self.sample_rate) / float(self.hop_size)

    @classmethod
    def from_config(cls, od_config: OmniDiffusionConfig) -> Cosmos3SoundTokenizer:
        args = _pipeline_args(od_config)
        model_path = getattr(od_config, "model", None)
        explicit_avae_path = (
            args.get("sound_tokenizer_path")
            or args.get("avae_path")
            or args.get("cosmos3_avae_path")
            or os.environ.get("COSMOS3_SOUND_TOKENIZER_PATH")
        )
        explicit_config_path = args.get("sound_tokenizer_config_path") or os.environ.get(
            "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH"
        )

        model_root = str(model_path) if model_path and os.path.isdir(model_path) else None
        if model_root is None and model_path and not explicit_avae_path:
            from huggingface_hub import snapshot_download

            model_root = snapshot_download(
                repo_id=str(model_path),
                revision=getattr(od_config, "revision", None),
                allow_patterns=[
                    f"{SOUND_TOKENIZER_COMPONENT_NAME}/config.json",
                    f"{SOUND_TOKENIZER_COMPONENT_NAME}/{SOUND_TOKENIZER_CHECKPOINT_NAME}",
                ],
            )

        if explicit_avae_path:
            avae_path = _resolve_model_file(explicit_avae_path, model_root)
        else:
            tokenizer_dir = Path(model_root) / SOUND_TOKENIZER_COMPONENT_NAME if model_root else None
            candidate = tokenizer_dir / SOUND_TOKENIZER_CHECKPOINT_NAME if tokenizer_dir else None
            avae_path = str(candidate) if candidate and candidate.exists() else None

        if not avae_path:
            raise ValueError(
                "Cosmos3 sound generation was requested, but no AVAE sound "
                "tokenizer checkpoint was provided. Set "
                "custom_pipeline_args['sound_tokenizer_path'] or "
                "COSMOS3_SOUND_TOKENIZER_PATH, or include "
                f"{SOUND_TOKENIZER_COMPONENT_NAME}/{SOUND_TOKENIZER_CHECKPOINT_NAME} under the model path."
            )

        config_path = _resolve_model_file(explicit_config_path, model_root)
        if config_path is None and model_root:
            candidate = Path(model_root) / SOUND_TOKENIZER_COMPONENT_NAME / "config.json"
            config_path = str(candidate) if candidate.exists() else None
        component_config = _load_sound_tokenizer_component_config(config_path)
        component_values = _component_arch_values(component_config)

        sample_rate = _resolve_arch_value(
            od_config,
            args,
            component_values,
            field="sample_rate",
            custom_keys=("sound_sample_rate", "sample_rate"),
            nested_keys=("sample_rate", "sampling_rate"),
            top_level_keys=("sound_sample_rate", "sample_rate"),
            default=DEFAULT_SOUND_SAMPLE_RATE,
            cast=int,
        )
        audio_channels = _resolve_arch_value(
            od_config,
            args,
            component_values,
            field="audio_channels",
            custom_keys=("sound_audio_channels", "audio_channels", "stereo"),
            nested_keys=("audio_channels", "dec_out_channels", "stereo"),
            top_level_keys=("sound_audio_channels", "audio_channels", "stereo"),
            default=DEFAULT_SOUND_CHANNELS,
            cast=_as_audio_channels,
        )
        sound_dim = _resolve_arch_value(
            od_config,
            args,
            component_values,
            field="io_channels",
            custom_keys=("sound_dim", "io_channels", "latent_ch"),
            nested_keys=("io_channels", "vocoder_input_dim", "latent_ch"),
            top_level_keys=("sound_dim",),
            default=DEFAULT_SOUND_DIM,
            cast=int,
        )
        hop_size = _resolve_arch_value(
            od_config,
            args,
            component_values,
            field="hop_size",
            custom_keys=("sound_hop_size", "hop_size"),
            nested_keys=("hop_size",),
            top_level_keys=("sound_hop_size", "hop_size"),
            default=DEFAULT_SOUND_HOP_SIZE,
            cast=int,
        )
        normalize_latents = _as_bool(
            _resolve_normalization_value(
                od_config,
                args,
                name="normalize_latents",
                default=DEFAULT_SOUND_NORMALIZE_LATENTS,
            )
        )
        normalization_type = str(
            _resolve_normalization_value(
                od_config,
                args,
                name="normalization_type",
                default=DEFAULT_SOUND_NORMALIZATION_TYPE,
            )
        )
        tanh_input_scale = float(
            _resolve_normalization_value(
                od_config,
                args,
                name="tanh_input_scale",
                default=DEFAULT_SOUND_TANH_INPUT_SCALE,
            )
        )
        tanh_output_scale = float(
            _resolve_normalization_value(
                od_config,
                args,
                name="tanh_output_scale",
                default=DEFAULT_SOUND_TANH_OUTPUT_SCALE,
            )
        )
        tanh_clamp = float(
            _resolve_normalization_value(
                od_config,
                args,
                name="tanh_clamp",
                default=DEFAULT_SOUND_TANH_CLAMP,
            )
        )
        tokenizer = Cosmos3AVAEAudioTokenizer(
            checkpoint_path=str(avae_path),
            config_path=config_path,
            sample_rate=sample_rate,
            audio_channels=audio_channels,
            io_channels=sound_dim,
            hop_size=hop_size,
            normalize_latents=normalize_latents,
            normalization_type=normalization_type,
            tanh_input_scale=tanh_input_scale,
            tanh_output_scale=tanh_output_scale,
            tanh_clamp=tanh_clamp,
            dtype=getattr(od_config, "dtype", torch.bfloat16),
            device=get_local_device(),
        )
        if _is_rank_zero():
            logger.info(
                "Loaded Cosmos3 AVAE sound tokenizer from %s "
                "(sr=%d, channels=%d, latent_ch=%d, hop=%d, latent_fps=%.3f)",
                avae_path,
                sample_rate,
                audio_channels,
                sound_dim,
                hop_size,
                float(sample_rate) / float(hop_size),
            )
        return cls(tokenizer)

    def get_latent_num_samples(self, num_audio_samples: int) -> int:
        return int(self.tokenizer.get_latent_num_samples(num_audio_samples))

    def get_audio_num_samples(self, num_latent_samples: int) -> int:
        return int(self.tokenizer.get_audio_num_samples(num_latent_samples))

    @torch.no_grad()
    def decode(self, latents: torch.Tensor) -> torch.Tensor:
        """Decode sound latents.

        Args:
            latents: ``[B, C, T]`` or ``[C, T]`` tensor.

        Returns:
            ``[B, audio_channels, N]`` tensor for batched input, or
            ``[audio_channels, N]`` for unbatched input.
        """
        squeeze = latents.ndim == 2
        if squeeze:
            latents = latents.unsqueeze(0)
        audio = self.tokenizer.decode(latents)
        audio = audio.clamp(-1.0, 1.0)
        return audio.squeeze(0) if squeeze else audio