# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Omni-specific argument parsing for python -m dynamo.vllm.omni."""

import argparse
import logging
from typing import Optional

from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs

try:
    from vllm.utils import FlexibleArgumentParser
except ImportError:
    from vllm.utils.argparse_utils import FlexibleArgumentParser

from dynamo.common.configuration.arg_group import ArgGroup
from dynamo.common.configuration.groups.runtime_args import (
    DynamoRuntimeArgGroup,
    DynamoRuntimeConfig,
)
from dynamo.common.configuration.utils import add_argument, add_negatable_bool_argument

logger = logging.getLogger(__name__)


class OmniArgGroup(ArgGroup):
    """Diffusion pipeline kwargs passed through to AsyncOmni() constructor.

    These are NOT part of AsyncOmniEngineArgs (which handles vLLM engine-level
    args like model, tp, max_model_len). Instead they are direct constructor
    kwargs for AsyncOmni and need Dynamo-side env-var (DYN_OMNI_*) support,
    so we define them here rather than relying on the upstream arg parser.
    """

    name = "dynamo-omni"

    def add_arguments(self, parser) -> None:
        g = parser.add_argument_group(
            "Omni Diffusion Options",
            "Diffusion pipeline parameters for vLLM-Omni multi-stage generation.",
        )

        add_argument(
            g,
            flag_name="--stage-configs-path",
            env_var="DYN_OMNI_STAGE_CONFIGS_PATH",
            default=None,
            help="Path to vLLM-Omni stage configuration YAML file (optional).",
        )

        # Video encoding
        add_argument(
            g,
            flag_name="--default-video-fps",
            env_var="DYN_OMNI_DEFAULT_VIDEO_FPS",
            default=16,
            arg_type=int,
            help="Default frames per second for generated videos.",
        )

        # Layerwise offloading
        add_negatable_bool_argument(
            g,
            flag_name="--enable-layerwise-offload",
            env_var="DYN_OMNI_ENABLE_LAYERWISE_OFFLOAD",
            default=False,
            help="Enable layerwise (blockwise) offloading on DiT modules to reduce GPU memory.",
        )
        add_argument(
            g,
            flag_name="--layerwise-num-gpu-layers",
            env_var="DYN_OMNI_LAYERWISE_NUM_GPU_LAYERS",
            default=1,
            arg_type=int,
            help="Number of ready layers (blocks) to keep on GPU during generation.",
        )

        # VAE optimization
        add_negatable_bool_argument(
            g,
            flag_name="--vae-use-slicing",
            env_var="DYN_OMNI_VAE_USE_SLICING",
            default=False,
            help="Enable VAE slicing for memory optimization in diffusion models.",
        )
        add_negatable_bool_argument(
            g,
            flag_name="--vae-use-tiling",
            env_var="DYN_OMNI_VAE_USE_TILING",
            default=False,
            help="Enable VAE tiling for memory optimization in diffusion models.",
        )

        # Diffusion scheduling
        add_argument(
            g,
            flag_name="--boundary-ratio",
            env_var="DYN_OMNI_BOUNDARY_RATIO",
            default=0.875,
            arg_type=float,
            help=(
                "Boundary split ratio for low/high DiT transformers. "
                "Default 0.875 uses both transformers for best quality. "
                "Set to 1.0 to load only the low-noise transformer (saves memory)."
            ),
        )
        add_argument(
            g,
            flag_name="--flow-shift",
            env_var="DYN_OMNI_FLOW_SHIFT",
            default=None,
            arg_type=float,
            help="Scheduler flow_shift parameter (5.0 for 720p, 12.0 for 480p).",
        )

        # Cache acceleration
        add_argument(
            g,
            flag_name="--cache-backend",
            env_var="DYN_OMNI_CACHE_BACKEND",
            default=None,
            choices=["cache_dit", "tea_cache"],
            help=(
                "Cache backend for diffusion acceleration. "
                "'cache_dit' enables DBCache + SCM + TaylorSeer. "
                "'tea_cache' enables TeaCache."
            ),
        )
        add_argument(
            g,
            flag_name="--cache-config",
            env_var="DYN_OMNI_CACHE_CONFIG",
            default=None,
            help="Cache configuration as JSON string (overrides defaults).",
        )
        add_negatable_bool_argument(
            g,
            flag_name="--enable-cache-dit-summary",
            env_var="DYN_OMNI_ENABLE_CACHE_DIT_SUMMARY",
            default=False,
            help="Enable cache-dit summary logging after diffusion forward passes.",
        )

        # Execution mode
        add_negatable_bool_argument(
            g,
            flag_name="--enable-cpu-offload",
            env_var="DYN_OMNI_ENABLE_CPU_OFFLOAD",
            default=False,
            help="Enable CPU offloading for diffusion models to reduce GPU memory usage.",
        )
        add_negatable_bool_argument(
            g,
            flag_name="--enforce-eager",
            env_var="DYN_OMNI_ENFORCE_EAGER",
            default=False,
            help="Disable torch.compile and force eager execution for diffusion models.",
        )

        # Diffusion parallel configuration
        add_argument(
            g,
            flag_name="--ulysses-degree",
            env_var="DYN_OMNI_ULYSSES_DEGREE",
            default=1,
            arg_type=int,
            help="Number of GPUs used for Ulysses sequence parallelism in diffusion.",
        )
        add_argument(
            g,
            flag_name="--ring-degree",
            env_var="DYN_OMNI_RING_DEGREE",
            default=1,
            arg_type=int,
            help="Number of GPUs used for ring sequence parallelism in diffusion.",
        )
        add_argument(
            g,
            flag_name="--cfg-parallel-size",
            env_var="DYN_OMNI_CFG_PARALLEL_SIZE",
            default=1,
            arg_type=int,
            choices=[1, 2],
            help="Number of GPUs used for classifier free guidance parallelism.",
        )


class OmniConfig(DynamoRuntimeConfig):
    """Configuration for Dynamo vLLM-Omni worker."""

    component: str = "backend"
    endpoint: Optional[str] = None

    # mirror vLLM
    model: str
    served_model_name: Optional[str] = None

    # vLLM-Omni engine args
    engine_args: AsyncOmniEngineArgs

    # OmniArgGroup fields (populated by from_cli_args)
    stage_configs_path: Optional[str] = None
    default_video_fps: int = 16
    enable_layerwise_offload: bool = False
    layerwise_num_gpu_layers: int = 1
    vae_use_slicing: bool = False
    vae_use_tiling: bool = False
    boundary_ratio: float = 0.875
    flow_shift: Optional[float] = None
    cache_backend: Optional[str] = None
    cache_config: Optional[str] = None
    enable_cache_dit_summary: bool = False
    enable_cpu_offload: bool = False
    enforce_eager: bool = False
    ulysses_degree: int = 1
    ring_degree: int = 1
    cfg_parallel_size: int = 1

    def validate(self) -> None:
        DynamoRuntimeConfig.validate(self)
        if self.default_video_fps <= 0:
            raise ValueError("--default-video-fps must be > 0")
        if self.ulysses_degree <= 0:
            raise ValueError("--ulysses-degree must be > 0")
        if self.ring_degree <= 0:
            raise ValueError("--ring-degree must be > 0")
        if not (0 < self.boundary_ratio <= 1):
            raise ValueError("--boundary-ratio must be in (0, 1]")


def parse_omni_args() -> OmniConfig:
    """Parse command-line arguments for the vLLM-Omni backend."""
    dynamo_runtime_argspec = DynamoRuntimeArgGroup()
    omni_argspec = OmniArgGroup()

    parser = argparse.ArgumentParser(
        description="Dynamo vLLM-Omni worker",
        formatter_class=argparse.RawTextHelpFormatter,
        allow_abbrev=False,
    )

    dynamo_runtime_argspec.add_arguments(parser)
    omni_argspec.add_arguments(parser)

    # Add vLLM-Omni engine args
    vg = parser.add_argument_group(
        "vLLM-Omni Engine Options. Please refer to vLLM-Omni documentation for more details."
    )
    vllm_parser = FlexibleArgumentParser(add_help=False)
    AsyncOmniEngineArgs.add_cli_args(vllm_parser, async_args_only=False)

    for action in vllm_parser._actions:
        if not action.option_strings:
            continue
        vg._group_actions.append(action)

    args, unknown = parser.parse_known_args()
    config = OmniConfig.from_cli_args(args)

    # Default endpoint to "generate" if not explicitly provided by user
    if config.endpoint is None:
        config.endpoint = "generate"

    vllm_args = vllm_parser.parse_args(unknown)
    config.model = vllm_args.model

    engine_args = AsyncOmniEngineArgs.from_cli_args(vllm_args)

    if getattr(engine_args, "served_model_name", None) is not None:
        served = engine_args.served_model_name
        if len(served) > 1:
            raise ValueError("We do not support multiple model names.")
        config.served_model_name = served[0]

    config.engine_args = engine_args
    config.validate()
    return config
