# SPDX-License-Identifier: Apache-2.0

import argparse
import logging
import os

from sglang.cli.utils import get_is_diffusion_model, get_model_path
from sglang.srt.utils import kill_process_tree
from sglang.srt.utils.common import suppress_noisy_warnings

suppress_noisy_warnings()

logger = logging.getLogger(__name__)


def _extract_model_type_override(extra_argv):
    """Extract and remove --model-type override from argv."""
    model_type = "auto"
    filtered_argv = []
    i = 0
    while i < len(extra_argv):
        arg = extra_argv[i]
        if arg == "--model-type":
            if i + 1 >= len(extra_argv):
                raise Exception(
                    "Error: --model-type requires a value. "
                    "Valid values are: auto, llm, diffusion."
                )
            model_type = extra_argv[i + 1]
            i += 2
            continue

        if arg.startswith("--model-type="):
            model_type = arg.split("=", 1)[1]
            i += 1
            continue

        filtered_argv.append(arg)
        i += 1

    if model_type not in ("auto", "llm", "diffusion"):
        raise Exception(
            f"Error: invalid --model-type '{model_type}'. "
            "Valid values are: auto, llm, diffusion."
        )
    return model_type, filtered_argv


def _normalize_positional_model_path(extra_argv):
    """Allow `sglang serve <model>` while preserving existing flag parsing."""
    if extra_argv and not extra_argv[0].startswith("-"):
        return ["--model-path", extra_argv[0], *extra_argv[1:]], True
    return extra_argv, False


def serve(args, extra_argv):
    if any(h in extra_argv for h in ("-h", "--help")):
        # Since the server type is determined by the model, and we don't have a model path,
        # we can't show the exact help. Instead, we show a general help message and then
        # the help for both possible server types.
        print(
            "Usage: sglang serve <model-name-or-path> [additional-arguments]\n"
            "   or: sglang serve --model-path <model-name-or-path> [additional-arguments]\n\n"
            "This command can launch either a standard language model server or a diffusion model server.\n"
            "The server type is determined by the --model-path.\n"
            "Optional override: --model-type {auto,llm,diffusion} "
            "(default: auto, fallback to LLM on detection failure)."
        )

        print("\n--- Help for Standard Language Model Server ---")
        from sglang.srt.server_args import prepare_server_args

        try:
            prepare_server_args(["--help"])
        except SystemExit:
            pass  # argparse --help calls sys.exit

        print("\n--- Help for Diffusion Model Server ---")
        try:
            from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
                add_multimodal_gen_serve_args,
            )

            parser = argparse.ArgumentParser(
                prog="sglang serve",
                description="SGLang Diffusion Model Serving",
            )
            add_multimodal_gen_serve_args(parser)
            parser.print_help()
        except ImportError:
            print(
                "Diffusion model support is not available. "
                'Install with: pip install "sglang[diffusion]"'
            )
        return

    from sglang.srt.plugins import load_plugins

    load_plugins()

    model_type, dispatch_argv = _extract_model_type_override(extra_argv)
    dispatch_argv, positional_model_path = _normalize_positional_model_path(
        dispatch_argv
    )
    model_path = get_model_path(dispatch_argv)
    try:
        if model_type == "auto":
            is_diffusion_model = get_is_diffusion_model(model_path)
            if is_diffusion_model:
                logger.info("Diffusion model detected")
        else:
            is_diffusion_model = model_type == "diffusion"
            logger.info(
                "Dispatch override enabled: --model-type=%s " "(skip auto detection)",
                model_type,
            )

        if is_diffusion_model:
            # Logic for Diffusion Models
            from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
                add_multimodal_gen_serve_args,
                execute_serve_cmd,
            )

            parser = argparse.ArgumentParser(
                description="SGLang Diffusion Model Serving"
            )
            add_multimodal_gen_serve_args(parser)
            parsed_args, remaining_argv = parser.parse_known_args(dispatch_argv)
            if positional_model_path:
                parsed_args._sglang_explicit_arg_names = {"model_path"}

            execute_serve_cmd(parsed_args, remaining_argv)
        else:
            # Logic for Standard Language Models
            from sglang.launch_server import run_server
            from sglang.srt.server_args import prepare_server_args

            server_args = prepare_server_args(dispatch_argv)

            run_server(server_args)
    finally:
        kill_process_tree(os.getpid(), include_parent=False)