# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
import os
import subprocess
import sys
import tempfile
import traceback

import pytest
import torch

# Cache for available backends detection (computed once at import time)
_AVAILABLE_BACKENDS = None


def get_available_backends():
    """Detect which backends are available in the current environment.

    Returns:
        set: A set of available backend names ('trtllm', 'vllm', 'sglang')
    """
    global _AVAILABLE_BACKENDS
    if _AVAILABLE_BACKENDS is not None:
        return _AVAILABLE_BACKENDS

    available = set()

    try:
        import tensorrt_llm  # noqa: F401

        available.add("trtllm")
    except ImportError:
        pass

    try:
        import vllm  # noqa: F401

        available.add("vllm")
    except ImportError:
        pass

    try:
        import sglang  # noqa: F401

        available.add("sglang")
    except ImportError:
        pass

    _AVAILABLE_BACKENDS = available
    print(f"[deploy_utils] Detected available backends: {available}")
    return _AVAILABLE_BACKENDS


def _run_trtllm_deploy(
    model_id: str,
    tensor_parallel_size: int,
    mini_sm: int,
    attn_backend: str,
    base_model: str,
    eagle3_one_model: bool,
    engine_kwargs: dict | None = None,
) -> None:
    """Top-level entry for subprocess: run TensorRT-LLM deploy in a child process."""
    try:
        deployer = ModelDeployer(
            backend="trtllm",
            model_id=model_id,
            tensor_parallel_size=tensor_parallel_size,
            mini_sm=mini_sm,
            attn_backend=attn_backend,
            base_model=base_model,
            eagle3_one_model=eagle3_one_model,
            engine_kwargs=engine_kwargs,
        )
        deployer._deploy_trtllm_impl()
    except Exception:
        traceback.print_exc()
        pytest.fail(traceback.format_exc())


def _run_vllm_deploy(
    model_id: str,
    tensor_parallel_size: int,
    mini_sm: int,
    attn_backend: str,
    base_model: str,
    eagle3_one_model: bool,
    engine_kwargs: dict | None = None,
) -> None:
    """Top-level entry for subprocess: run vLLM deploy in a child process."""
    try:
        deployer = ModelDeployer(
            backend="vllm",
            model_id=model_id,
            tensor_parallel_size=tensor_parallel_size,
            mini_sm=mini_sm,
            attn_backend=attn_backend,
            base_model=base_model,
            eagle3_one_model=eagle3_one_model,
            engine_kwargs=engine_kwargs,
        )
        deployer._deploy_vllm_impl()
    except Exception:
        traceback.print_exc()
        pytest.fail(traceback.format_exc())


def _run_sglang_deploy(
    model_id: str,
    tensor_parallel_size: int,
    mini_sm: int,
    attn_backend: str,
    base_model: str,
    eagle3_one_model: bool,
    engine_kwargs: dict | None = None,
) -> None:
    """Top-level entry for subprocess: run SGLang deploy in a child process."""
    try:
        deployer = ModelDeployer(
            backend="sglang",
            model_id=model_id,
            tensor_parallel_size=tensor_parallel_size,
            mini_sm=mini_sm,
            attn_backend=attn_backend,
            base_model=base_model,
            eagle3_one_model=eagle3_one_model,
            engine_kwargs=engine_kwargs,
        )
        deployer._deploy_sglang_impl()
    except Exception:
        traceback.print_exc()
        pytest.fail(traceback.format_exc())


def _run_deploy_via_subprocess(
    backend: str,
    model_id: str,
    tensor_parallel_size: int,
    mini_sm: int,
    attn_backend: str,
    base_model: str,
    eagle3_one_model: bool,
    engine_kwargs: dict,
) -> None:
    """Run deploy in a subprocess and print its stdout/stderr so pytest capture=tee-sys captures to DB."""
    tests_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    project_root = os.path.dirname(tests_dir)
    env = {
        **os.environ,
        "PYTHONPATH": project_root + os.pathsep + os.environ.get("PYTHONPATH", ""),
    }
    code = f"""import sys
sys.path.insert(0, {tests_dir!r})
if __name__ == '__main__':
    from _test_utils.deploy_utils import _run_{backend}_deploy
    _run_{backend}_deploy(
        {model_id!r}, {tensor_parallel_size}, {mini_sm}, {attn_backend!r}, {base_model!r},
        {eagle3_one_model}, {engine_kwargs!r}
    )
"""
    if backend == "trtllm":
        mpirun = os.environ.get("MPIRUN", "mpirun")
        cmd = [
            mpirun,
            "--allow-run-as-root",
            "--oversubscribe",
            "-n",
            "1",
            sys.executable,
        ]
    else:
        cmd = [sys.executable, "-c", code]

    if backend == "trtllm":
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write(code)
            tmp_path = f.name
        try:
            result = subprocess.run(
                [*cmd, tmp_path],
                cwd=tests_dir,
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
            )
        finally:
            os.unlink(tmp_path)
    else:
        result = subprocess.run(
            cmd,
            cwd=tests_dir,
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
    if result.stdout:
        print(result.stdout, end="", flush=True)
    result.check_returncode()


# Common test prompts for all backends
COMMON_PROMPTS = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]


class ModelDeployer:
    def __init__(
        self,
        backend: str = "trtllm",
        model_id: str = "",
        tensor_parallel_size: int = 1,
        mini_sm: int = 89,
        attn_backend: str = "TRTLLM",
        base_model: str = "",
        eagle3_one_model: bool = True,
        engine_kwargs: dict | None = None,
    ):
        """
        Initialize the ModelDeployer.

        Args:
            backend: The backend to use ('vllm', 'trtllm', or 'sglang')
            model_id: Path to the model
            tensor_parallel_size: Tensor parallel size for distributed inference
            mini_sm: Minimum SM (Streaming Multiprocessor) requirement for the model
            attn_backend: is for TRT LLM deployment
            engine_kwargs: extra keyword arguments for this backend's engine constructor,
                overriding the defaults below. For a checkpoint whose model card asks for
                settings the defaults do not cover, pass them here rather than special-casing
                the model id.
        """
        self.backend = backend
        self.model_id = model_id
        self.tensor_parallel_size = tensor_parallel_size
        self.mini_sm = mini_sm
        self.attn_backend = attn_backend
        self.base_model = base_model
        self.eagle3_one_model = eagle3_one_model
        self.engine_kwargs = dict(engine_kwargs or {})

    def run(self):
        """Run the deployment based on the specified backend."""
        if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
            pytest.skip("CUDA is not available")
            return
        if torch.cuda.get_device_capability() < (
            self.mini_sm // 10,
            self.mini_sm % 10,
        ):
            pytest.skip(reason=f"Requires sm{self.mini_sm} or higher")
            return

        if torch.cuda.device_count() < self.tensor_parallel_size:
            pytest.skip(reason=f"Requires at least {self.tensor_parallel_size} GPUs")
            return

        print(f"Deploying model: {self.model_id} with backend: {self.backend}")
        print(f"Tensor parallel size: {self.tensor_parallel_size}")
        # Use subprocess + capture so pytest capture=tee-sys (and DB plugins) see deploy output.
        if self.backend in ("vllm", "trtllm", "sglang"):
            _run_deploy_via_subprocess(
                backend=self.backend,
                model_id=self.model_id,
                tensor_parallel_size=self.tensor_parallel_size,
                mini_sm=self.mini_sm,
                attn_backend=self.attn_backend,
                base_model=self.base_model,
                eagle3_one_model=self.eagle3_one_model,
                engine_kwargs=self.engine_kwargs,
            )
        else:
            raise ValueError(f"Unknown backend: {self.backend}")

    def _deploy_trtllm_impl(self):
        """Run TensorRT-LLM deploy (used by subprocess in run())."""
        from tensorrt_llm import LLM, SamplingParams
        from tensorrt_llm.llmapi import CudaGraphConfig, EagleDecodingConfig, KvCacheConfig

        sampling_params = SamplingParams(max_tokens=32)
        qwen3_models = (
            "nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4",
            "nvidia/Qwen3-Next-80B-A3B-Thinking-NVFP4",
        )
        nemotron_models = (
            "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",
            "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
        )
        kv_cache_config = KvCacheConfig(
            enable_block_reuse=self.model_id not in qwen3_models,
            free_gpu_memory_fraction=0.8,
            mamba_ssm_cache_dtype="float32" if self.model_id not in nemotron_models else "auto",
        )
        base_kw = {
            "tensor_parallel_size": self.tensor_parallel_size,
            "enable_attention_dp": False,
            "attn_backend": self.attn_backend,
            "trust_remote_code": True,
            "max_batch_size": 8,
            **self.engine_kwargs,
        }

        if "eagle" in self.model_id.lower():
            spec_config = EagleDecodingConfig(
                max_draft_len=3,
                speculative_model_dir=self.model_id,
                eagle3_one_model=self.eagle3_one_model,
            )
            llm = LLM(
                model=self.base_model,
                disable_overlap_scheduler=True,
                enable_autotuner=False,
                speculative_config=spec_config,
                cuda_graph_config=CudaGraphConfig(max_batch_size=1),
                kv_cache_config=kv_cache_config,
                **base_kw,
            )
        else:
            llm = LLM(
                model=self.model_id,
                backend="pytorch",
                kv_cache_config=kv_cache_config,
                **base_kw,
            )

        # Deferred import: transformers is a heavy optional dependency not always present.
        from transformers import AutoTokenizer

        tokenizer_model = self.base_model if "eagle" in self.model_id.lower() else self.model_id
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_model, trust_remote_code=True)
        prompts = COMMON_PROMPTS
        if tokenizer.chat_template is not None:
            prompts = [
                tokenizer.apply_chat_template(
                    [{"role": "user", "content": prompt}],
                    tokenize=False,
                    add_generation_prompt=True,
                )
                for prompt in COMMON_PROMPTS
            ]

        outputs = llm.generate(prompts, sampling_params)
        assert len(outputs) == len(COMMON_PROMPTS), (
            f"Expected {len(COMMON_PROMPTS)} outputs, got {len(outputs)}"
        )
        for i, output in enumerate(outputs):
            generated_text = output.outputs[0].text
            assert isinstance(generated_text, str), f"Output {i} text is not a string"
            assert generated_text.strip(), f"Output {i} generated empty text"
            print(f"Model: {self.model_id}")
            print(f"Input prompt: {COMMON_PROMPTS[i]!r}")
            print(f"Rendered prompt: {output.prompt!r}")
            print(f"Generated text: {generated_text!r}")
            print("-" * 50)
        del llm

    def _deploy_vllm_impl(self):
        """Run vLLM deploy (used by subprocess in run())."""
        from vllm import LLM, SamplingParams

        if "eagle" in self.model_id.lower():
            llm = LLM(
                model=self.base_model,
                speculative_config={
                    "method": "eagle3",
                    "model": self.model_id,
                    "num_speculative_tokens": 3,
                },
                tensor_parallel_size=self.tensor_parallel_size,
                trust_remote_code=True,
                max_model_len=4096,
            )
        else:
            quantization_method = "modelopt"
            if "fp4" in self.model_id.lower():
                quantization_method = "modelopt_fp4"
            # DeepSeek-V4 FlashMLA requires fp8 kv-cache explicitly
            kv_cache_dtype = "fp8" if "deepseek-v4" in self.model_id.lower() else "auto"
            llm_kwargs = {
                "model": self.model_id,
                "quantization": quantization_method,
                "tensor_parallel_size": self.tensor_parallel_size,
                "trust_remote_code": True,
                "max_model_len": 4096,
                "kv_cache_dtype": kv_cache_dtype,
                **self.engine_kwargs,
            }
            llm = LLM(**llm_kwargs)
        sampling_params = SamplingParams(temperature=0.8, top_p=0.9)
        conversations = [[{"role": "user", "content": p}] for p in COMMON_PROMPTS]
        outputs = llm.chat(conversations, sampling_params)

        assert len(outputs) == len(COMMON_PROMPTS), (
            f"Expected {len(COMMON_PROMPTS)} outputs, got {len(outputs)}"
        )

        for i, output in enumerate(outputs):
            assert hasattr(output, "outputs"), f"Output {i} missing 'outputs' attribute"
            assert len(output.outputs) > 0, f"Output {i} has no generated text"
            assert hasattr(output.outputs[0], "text"), f"Output {i} missing 'text' attribute"
            assert isinstance(output.outputs[0].text, str), f"Output {i} text is not a string"
            generated_text = output.outputs[0].text
            assert generated_text.strip(), f"Output {i} generated empty text"

            print(f"Model: {self.model_id}")
            print(f"Input prompt: {COMMON_PROMPTS[i]!r}")
            print(f"Rendered prompt: {getattr(output, 'prompt', None)!r}")
            print(f"Generated text: {generated_text!r}")
            print("-" * 50)
        del llm

    def _deploy_sglang_impl(self):
        """Run SGLang deploy (used by subprocess in run())."""
        import sglang as sgl

        quantization_method = "modelopt"
        if "fp4" in self.model_id.lower():
            quantization_method = "modelopt_fp4"
        if "eagle" in self.model_id.lower():
            eagle_kwargs = {
                "model_path": self.base_model,
                "speculative_algorithm": "EAGLE3",
                "speculative_num_steps": 3,
                "speculative_eagle_topk": 1,
                "speculative_num_draft_tokens": 4,
                "speculative_draft_model_path": self.model_id,
                "tp_size": self.tensor_parallel_size,
                "trust_remote_code": True,
                "mem_fraction_static": 0.7,
                "context_length": 1024,
            }
            # Nemotron hybrid (Mamba+attention) requires extra_buffer scheduler
            # strategy and SGLANG_ENABLE_SPEC_V2 for radix cache compatibility
            # with speculative decoding
            if "nemotron" in self.base_model.lower():
                eagle_kwargs["mamba_scheduler_strategy"] = "extra_buffer"
                os.environ["SGLANG_ENABLE_SPEC_V2"] = "1"
            llm = sgl.Engine(**eagle_kwargs)
        elif self.model_id in (
            "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",
            "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",
            "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-FP8",
            "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4",
        ):
            llm = sgl.Engine(
                model_path=self.model_id,
                quantization=quantization_method,
                tp_size=self.tensor_parallel_size,
                trust_remote_code=True,
                attention_backend="flashinfer",
            )
        else:
            sgl_kwargs = {
                "model_path": self.model_id,
                "quantization": quantization_method,
                "tp_size": self.tensor_parallel_size,
                "trust_remote_code": True,
                "context_length": 4096,
                **self.engine_kwargs,
            }
            llm = sgl.Engine(**sgl_kwargs)
        # Deferred import: transformers is a heavy optional dependency not always present.
        from transformers import AutoTokenizer

        tokenizer_model = self.base_model if "eagle" in self.model_id.lower() else self.model_id
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_model, trust_remote_code=True)
        if tokenizer.chat_template is not None:
            prompts = [
                tokenizer.apply_chat_template(
                    [{"role": "user", "content": p}],
                    tokenize=False,
                    add_generation_prompt=True,
                )
                for p in COMMON_PROMPTS
            ]
        else:
            prompts = COMMON_PROMPTS

        outputs = llm.generate(prompts)
        assert len(outputs) == len(COMMON_PROMPTS), (
            f"Expected {len(COMMON_PROMPTS)} outputs, got {len(outputs)}"
        )
        for i, output in enumerate(outputs):
            if isinstance(output, dict):
                generated_text = output["text"]
            elif isinstance(output, str):
                generated_text = output
            else:
                raise TypeError(f"Output {i} has unexpected type {type(output)}: {output!r}")
            assert isinstance(generated_text, str), f"Output {i} text is not a string"
            assert generated_text.strip(), f"Output {i} generated empty text"
            print(f"Model: {self.model_id}")
            print(f"Input prompt: {COMMON_PROMPTS[i]!r}")
            print(f"Rendered prompt: {prompts[i]!r}")
            print(f"Generated text: {generated_text!r}")
            print("-" * 50)
        llm.shutdown()
        del llm


class ModelDeployerList:
    def __init__(self, **params):
        self.params = {}
        for key, value in params.items():
            if isinstance(value, (list, tuple)):
                self.params[key] = list(value)
            else:
                self.params[key] = [value]

        # Filter backends to only include available ones
        if "backend" in self.params:
            available = get_available_backends()
            original_backends = self.params["backend"]
            self.params["backend"] = [b for b in original_backends if b in available]

        # Pre-generate all deployers for pytest compatibility
        self._deployers = list(self._generate_deployers())

    def _generate_deployers(self):
        # If no backends available after filtering, yield nothing
        if "backend" in self.params and not self.params["backend"]:
            return

        for values in itertools.product(*self.params.values()):
            deployer = ModelDeployer(**dict(zip(self.params.keys(), values)))
            # Set test case ID in format "model_id_backend"
            deployer.test_id = f"{deployer.model_id}_{deployer.backend}"
            yield deployer

    def __iter__(self):
        return iter(self._deployers)

    def __len__(self):
        return len(self._deployers)

    def __getitem__(self, index):
        return self._deployers[index]

    def __str__(self):
        return f"ModelDeployerList({len(self._deployers)} items)"

    def __repr__(self):
        return f"ModelDeployerList({len(self._deployers)} items)"