# SPDX-License-Identifier: Apache-2.0
"""NVFP4Config import + lazy-flashinfer behavior tests.

The actual NVFP4 kernels need flashinfer + a CUDA device, so these
tests focus on the lazy-import contract: the module must load on
hosts without flashinfer, and only call sites should fail with a
clear error when flashinfer is missing.
"""
from __future__ import annotations

import sys
import types

import pytest


def test_nvfp4config_imports_without_flashinfer(monkeypatch):
    """Importing the module on a host without flashinfer must succeed.

    This is the contract Dreamverse depends on — the GPU worker boots
    even when flashinfer is not in the venv, because the import happens
    in `video_generation.py` before any NVFP4 kernel is invoked.
    """
    # Hide flashinfer from sys.modules and the import path.
    monkeypatch.setitem(sys.modules, "flashinfer", None)
    # Force a re-import of the target module.
    sys.modules.pop("fastvideo.layers.quantization.nvfp4_config", None)
    from fastvideo.layers.quantization.nvfp4_config import NVFP4Config
    config = NVFP4Config()
    assert config.get_name() == "nvfp4"
    assert config.layer_profile == "refine"


def test_nvfp4config_layer_profile_round_trips_from_dict():
    from fastvideo.layers.quantization.nvfp4_config import NVFP4Config
    config = NVFP4Config.from_config({"layer_profile": "base"})
    assert config.layer_profile == "base"
    config = NVFP4Config.from_config({})
    assert config.layer_profile == "refine"


def test_nvfp4_kernel_call_raises_clear_error_without_flashinfer(monkeypatch):
    """A call into the NVFP4 kernels must raise an actionable
    ImportError when flashinfer is missing, not a confusing
    AttributeError or NameError."""
    # Stage a fake flashinfer that fails on import.
    monkeypatch.setitem(sys.modules, "flashinfer", _raise_module_on_import("flashinfer"))
    sys.modules.pop("fastvideo.layers.quantization.nvfp4_config", None)
    from fastvideo.layers.quantization.nvfp4_config import _require_flashinfer
    with pytest.raises(ImportError, match="flashinfer"):
        _require_flashinfer()


def _raise_module_on_import(name: str) -> types.ModuleType:
    """Build a stub module that raises ImportError on any attribute
    access, so `from <name> import X` fails the way a missing package
    would."""

    class _RaisingModule(types.ModuleType):

        def __getattr__(self, item: str):
            raise ImportError(f"No module named '{name}.{item}'")

    return _RaisingModule(name)


def test_coerce_fp4_input_dtype_casts_and_rejects():
    """The FP4 input-dtype coercion (CPU-only, no flashinfer/CUDA): bf16/fp16
    pass through, other floats (the fp32 pre-attention norm in eager mode) are
    cast to bf16, and non-floating inputs are rejected fast."""
    import torch

    from fastvideo.layers.quantization.nvfp4_config import (_coerce_fp4_input_dtype)

    # bf16 / fp16 pass through untouched.
    bf16 = torch.zeros(4, 8, dtype=torch.bfloat16)
    assert _coerce_fp4_input_dtype(bf16) is bf16
    fp16 = torch.zeros(4, 8, dtype=torch.float16)
    assert _coerce_fp4_input_dtype(fp16) is fp16

    # Other floating dtypes (fp32 from an unfused norm, fp64) -> bf16.
    assert _coerce_fp4_input_dtype(torch.zeros(4, 8, dtype=torch.float32)).dtype is torch.bfloat16
    assert _coerce_fp4_input_dtype(torch.zeros(4, 8, dtype=torch.float64)).dtype is torch.bfloat16

    # Non-floating inputs are a real error, not silently cast.
    for bad in (torch.int32, torch.int64, torch.bool):
        with pytest.raises(TypeError, match="floating-point"):
            _coerce_fp4_input_dtype(torch.zeros(4, 8, dtype=bad))