import logging

import torch

from ltx_core.components.guiders import (
    MultiModalGuiderFactory,
    MultiModalGuiderParams,
    create_multimodal_guider_factory,
)
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.schedulers import LTX2Scheduler
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.loader.registry import Registry
from ltx_core.model.transformer import LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP, LTXAudioOnlyModelConfigurator
from ltx_core.model.transformer.compiling import CompilationConfig
from ltx_core.quantization import QuantizationPolicy
from ltx_core.types import Audio
from ltx_pipelines.utils import get_device
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
from ltx_pipelines.utils.args import (
    default_1_stage_t2a_arg_parser,
    resolve_cli_params,
)
from ltx_pipelines.utils.blocks import (
    AudioDecoder,
    DiffusionStage,
    PromptEncoder,
)
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
from ltx_pipelines.utils.media_io import encode_audio
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode

# Placeholder pixel dimensions used for ``VideoPixelShape`` construction.
# Audio-only generation reads ``frames`` and ``fps`` from the pixel shape via
# ``AudioLatentShape.from_video_pixel_shape`` (height/width are unused).
_AUDIO_ONLY_PLACEHOLDER_RES = 512


class T2AOneStagePipeline:
    """
    Single-stage text-to-audio generation pipeline.
    Generates audio at the target duration in a single diffusion pass with
    classifier-free guidance (CFG) on the audio modality only. The video
    modality is fully absent — the transformer runs audio-only by passing
    ``video=None`` to the ``DiffusionStage``.
    Assumes full non distilled model is provided in the checkpoint_path.
    """

    def __init__(
        self,
        checkpoint_path: str,
        gemma_root: str,
        loras: list[LoraPathStrengthAndSDOps],
        device: torch.device | None = None,
        quantization: QuantizationPolicy | None = None,
        registry: Registry | None = None,
        compilation_config: CompilationConfig | None = None,
        offload_mode: OffloadMode = OffloadMode.NONE,
        alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
    ):
        self.dtype = torch.bfloat16
        self.device = device or get_device()
        self._scheduler = LTX2Scheduler()
        self.prompt_encoder = PromptEncoder(
            checkpoint_path=checkpoint_path,
            gemma_root=gemma_root,
            dtype=self.dtype,
            device=self.device,
            registry=registry,
            offload_mode=offload_mode,
            alloc_trim_strategy=alloc_trim_strategy,
        )
        # Audio-only: build an audio-only transformer (model_configurator) so the video
        # weights are never instantiated, plus a use-case-specific SDOps that restricts
        # checkpoint reads to the audio model's keys, so the video weights are never even
        # read from disk (the loader skips any key the SDOps maps to None).
        self.stage = DiffusionStage.from_checkpoint(
            checkpoint_path=checkpoint_path,
            dtype=self.dtype,
            device=self.device,
            loras=tuple(loras),
            quantization=quantization,
            registry=registry,
            compilation_config=compilation_config,
            offload_mode=offload_mode,
            model_configurator=LTXAudioOnlyModelConfigurator,
            model_sd_ops=LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
            alloc_trim_strategy=alloc_trim_strategy,
        )
        self.audio_decoder = AudioDecoder(
            checkpoint_path=checkpoint_path,
            dtype=self.dtype,
            device=self.device,
            registry=registry,
            alloc_trim_strategy=alloc_trim_strategy,
        )

    def __call__(
        self,
        prompt: str,
        negative_prompt: str,
        seed: int,
        num_frames: int,
        frame_rate: float,
        num_inference_steps: int,
        audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
        enhance_prompt: bool = False,
        max_batch_size: int = 1,
        sigmas: torch.Tensor | None = None,
    ) -> Audio:
        generator = torch.Generator(device=self.device).manual_seed(seed)
        noiser = GaussianNoiser(generator=generator)

        ctx_p, ctx_n = self.prompt_encoder(
            [prompt, negative_prompt],
            enhance_first_prompt=enhance_prompt,
            enhance_prompt_image=None,
            enhance_prompt_seed=seed,
        )
        a_context_p = ctx_p.audio_encoding
        a_context_n = ctx_n.audio_encoding

        sigmas = (sigmas if sigmas is not None else self._scheduler.execute(steps=num_inference_steps)).to(
            dtype=torch.float32, device=self.device
        )

        # Normalize to a guider factory. Plain ``MultiModalGuiderParams`` (the default /
        # CLI case) becomes a simple sigma-independent guider, but callers may also pass
        # their own factory for sigma-dependent guidance; ``FactoryGuidedDenoiser`` always
        # consumes a factory.
        audio_guider_factory = create_multimodal_guider_factory(
            params=audio_guider_params,
            negative_context=a_context_n,
        )

        _, audio_state = self.stage(
            denoiser=FactoryGuidedDenoiser(
                v_context=None,
                a_context=a_context_p,
                video_guider_factory=None,
                audio_guider_factory=audio_guider_factory,
            ),
            sigmas=sigmas,
            noiser=noiser,
            width=_AUDIO_ONLY_PLACEHOLDER_RES,
            height=_AUDIO_ONLY_PLACEHOLDER_RES,
            frames=num_frames,
            fps=frame_rate,
            video=None,
            audio=ModalitySpec(context=a_context_p),
            max_batch_size=max_batch_size,
        )

        return self.audio_decoder(audio_state.latent)


@torch.inference_mode()
def main() -> None:
    logging.basicConfig(level=logging.INFO)
    params = resolve_cli_params()
    parser = default_1_stage_t2a_arg_parser(params=params)
    args = parser.parse_args()
    pipeline = T2AOneStagePipeline(
        checkpoint_path=args.checkpoint_path,
        gemma_root=args.gemma_root,
        loras=tuple(args.lora) if args.lora else (),
        quantization=args.quantization,
        compilation_config=args.compile,
        offload_mode=args.offload_mode,
    )
    audio = pipeline(
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        seed=args.seed,
        num_frames=args.num_frames,
        frame_rate=args.frame_rate,
        num_inference_steps=args.num_inference_steps,
        audio_guider_params=MultiModalGuiderParams(
            cfg_scale=args.audio_cfg_guidance_scale,
            stg_scale=args.audio_stg_guidance_scale,
            rescale_scale=args.audio_rescale_scale,
            # Audio-only generation has no video modality, so the video->audio
            # (v2a) cross-modal guidance is meaningless here. 1.0 disables it.
            modality_scale=1.0,
            skip_step=args.audio_skip_step,
            stg_blocks=args.audio_stg_blocks,
        ),
        max_batch_size=args.max_batch_size,
    )

    encode_audio(audio=audio, output_path=args.output_path)


if __name__ == "__main__":
    main()
