# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CPU regressions for Quack compilation and warmup."""

import sys
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch

from vllm_omni.quantization import quack_fp8

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


def test_compile_pool_checks_daemon_at_tuning_time(monkeypatch):
    process = SimpleNamespace(daemon=False)
    compiler_pool = object()
    pool = SimpleNamespace(
        pool_scope=Mock(spec=[], side_effect=lambda: nullcontext(compiler_pool)),
        suppress_pool=Mock(side_effect=lambda: nullcontext(object())),
    )
    original_scope = pool.pool_scope
    monkeypatch.setattr(quack_fp8, "current_process", lambda: process)
    monkeypatch.setattr(quack_fp8, "import_module", lambda name: pool)
    quack_fp8._configure_quack_compilation()
    with pool.pool_scope() as active_pool:
        assert active_pool is compiler_pool
    pool.suppress_pool.assert_not_called()

    # spawn installs the child's daemon flag after importing the worker module.
    process.daemon = True
    with pool.pool_scope() as active_pool:
        assert active_pool is None
    original_scope.assert_called_once_with()
    pool.suppress_pool.assert_called_once_with()


@pytest.fixture
def quack_dispatch(monkeypatch):
    """Exercise the installed patch with CPU substitutes for the GPU kernels."""

    def gemm(a, b, *, out, bias, alpha, tuned):
        assert isinstance(alpha, torch.Tensor)
        assert alpha.shape == (1,)
        assert alpha.dtype == torch.float32
        assert tuned is True
        result = (a.float() @ b.float()) * alpha
        if bias is not None:
            result += bias
        out.copy_(result)

    def flashinfer(a, b, *, out_dtype, scale_a, scale_b, bias):
        out = torch.full((a.shape[0], b.shape[1]), -7, dtype=out_dtype, device=a.device)
        return out + bias if bias is not None else out

    class FlashInferKernel:
        def apply_scaled_mm(self, **kwargs):
            raise AssertionError("The installed custom op should handle dispatch")

    gemm_mock = Mock(side_effect=gemm)
    fallback_mock = Mock(side_effect=flashinfer)
    monkeypatch.setattr(quack_fp8, "quack_enabled", lambda: True)
    monkeypatch.setattr(quack_fp8, "_gemm_interface", SimpleNamespace(gemm=gemm_mock))
    monkeypatch.setattr(quack_fp8, "_valid_scale_ptrs", set())
    monkeypatch.setattr(quack_fp8, "logger", Mock())
    monkeypatch.setitem(
        sys.modules,
        "vllm.model_executor.kernels.linear.scaled_mm.flashinfer",
        SimpleNamespace(FlashInferFP8ScaledMMLinearKernel=FlashInferKernel),
    )
    monkeypatch.setitem(sys.modules, "vllm.utils.flashinfer", SimpleNamespace(flashinfer_scaled_fp8_mm=fallback_mock))
    quack_fp8.install_quack_fp8_patch()
    kernel = FlashInferKernel()

    def forward(a, b, scale_a, scale_b, bias=None, out_dtype=torch.bfloat16):
        return kernel.apply_scaled_mm(
            A=a,
            B=b,
            As=scale_a,
            Bs=scale_b,
            out_dtype=out_dtype,
            bias=bias,
            output_shape=(2, a.shape[0] // 2, b.shape[1]),
        )

    return SimpleNamespace(forward=forward, gemm=gemm_mock, fallback=fallback_mock)


@pytest.fixture
def compile_fullgraph():
    torch._dynamo.reset()
    graphs = []

    def backend(graph, example_inputs):
        graphs.append(graph)
        return graph.forward

    def compile_fn(fn, **kwargs):
        return torch.compile(fn, backend=backend, fullgraph=True, **kwargs)

    try:
        yield compile_fn, graphs
    finally:
        torch._dynamo.reset()


def _fp8_inputs(m=4):
    a = torch.arange(m * 8, dtype=torch.float32).reshape(m, 8).to(torch.float8_e4m3fn)
    b = torch.ones(6, 8, dtype=torch.float8_e4m3fn).t()
    return a, b


@pytest.mark.parametrize("out_dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("with_bias", [False, True])
def test_scale_validation_does_not_recompile_for_each_layer(quack_dispatch, compile_fullgraph, out_dtype, with_bias):
    compile_fn, graphs = compile_fullgraph
    compiled = compile_fn(quack_dispatch.forward)
    with torch.inference_mode():
        a, b = _fp8_inputs()
        bias = torch.arange(6, dtype=out_dtype) if with_bias else None
        # Keep all pairs alive, as separate model layers do. Both pointer
        # addresses and scale values differ; tensor metadata stays identical.
        pairs = [(torch.tensor([float(i + 1)]), torch.tensor([0.5])) for i in range(5)]
        for _ in range(2):
            for scale_a, scale_b in pairs:
                expected = (a.float() @ b.float()) * scale_a * scale_b
                if bias is not None:
                    expected += bias
                out = compiled(a, b, scale_a, scale_b, bias, out_dtype)
                torch.testing.assert_close(out, expected.to(out_dtype).view(2, 2, 6))
    assert len(graphs) == 1
    assert any(node.target == torch.ops.vllm_omni.quack_fp8_scaled_mm.default for node in graphs[0].graph.nodes)
    assert quack_dispatch.gemm.call_count == 10
    quack_dispatch.fallback.assert_not_called()


def test_compiled_dispatch_rechecks_invalid_scales_after_loading(quack_dispatch, compile_fullgraph):
    compile_fn, graphs = compile_fullgraph
    compiled = compile_fn(quack_dispatch.forward)
    with torch.inference_mode():
        a, b = _fp8_inputs()
        pairs = []
        for bad_scale in (torch.finfo(torch.float32).min, 0.0, -1.0, float("nan"), float("inf")):
            for scale_index in (0, 1):
                scales = [torch.ones(1), torch.ones(1)]
                scales[scale_index].fill_(bad_scale)
                pairs.append(scales)
        for scale_a, scale_b in pairs:
            out = compiled(a, b, scale_a, scale_b)
            torch.testing.assert_close(out, torch.full((2, 2, 6), -7, dtype=torch.bfloat16))
        quack_dispatch.gemm.assert_not_called()
        assert quack_dispatch.fallback.call_count == len(pairs)

        # Populate the same buffers and switch branches without recompiling.
        for scale_a, scale_b in pairs:
            scale_a.fill_(0.25)
            scale_b.fill_(0.5)
            out = compiled(a, b, scale_a, scale_b)
            expected = ((a.float() @ b.float()) * 0.125).to(torch.bfloat16)
            torch.testing.assert_close(out, expected.view(2, 2, 6))
        assert quack_dispatch.gemm.call_count == len(pairs)
        assert quack_dispatch.fallback.call_count == len(pairs)
    assert len(graphs) == 1


def test_compiled_dispatch_falls_back_on_quack_failure(monkeypatch, quack_dispatch, compile_fullgraph):
    compile_fn, graphs = compile_fullgraph
    compiled = compile_fn(quack_dispatch.forward)
    with torch.inference_mode():
        a, b = _fp8_inputs()
        scale_a, scale_b = torch.ones(1), torch.ones(1)
        bias = torch.arange(6, dtype=torch.float32)
        expected = (torch.full((2, 2, 6), -7, dtype=torch.bfloat16) + bias).to(torch.bfloat16)
        compiled(a, b, scale_a, scale_b, bias)
        quack_dispatch.fallback.assert_not_called()

        quack_dispatch.gemm.side_effect = RuntimeError("kernel failed")
        out = compiled(a, b, scale_a, scale_b, bias)
        torch.testing.assert_close(out, expected)
        quack_fp8.logger.warning_once.assert_called_once()

        # The helper can also return None when Quack is unavailable.
        monkeypatch.setattr(quack_fp8, "_gemm_interface", False)
        out = compiled(a, b, scale_a, scale_b, bias)
        torch.testing.assert_close(out, expected)
        assert quack_dispatch.fallback.call_count == 2
    assert len(graphs) == 1


@pytest.mark.parametrize("scale_index", [0, 1])
def test_compiled_dispatch_falls_back_for_non_scalar_scales(quack_dispatch, compile_fullgraph, scale_index):
    compile_fn, graphs = compile_fullgraph
    compiled = compile_fn(quack_dispatch.forward)
    with torch.inference_mode():
        a, b = _fp8_inputs()
        scales = [torch.ones(1), torch.ones(1)]
        scales[scale_index] = torch.ones(2)
        out = compiled(a, b, *scales)
        torch.testing.assert_close(out, torch.full((2, 2, 6), -7, dtype=torch.bfloat16))
    quack_dispatch.gemm.assert_not_called()
    quack_dispatch.fallback.assert_called_once()
    assert len(graphs) == 1


@pytest.mark.parametrize("valid_scales", [False, True])
@pytest.mark.parametrize("with_bias", [False, True])
def test_dispatch_custom_op_contract(quack_dispatch, valid_scales, with_bias):
    a, b = _fp8_inputs()
    # opcheck's CPU mutation checker uses allclose, which does not support FP8.
    # The compile tests above exercise real FP8 inputs with these same mocks.
    a, b = a.float(), b.float()
    scale_a = torch.tensor([0.25 if valid_scales else torch.finfo(torch.float32).min])
    scale_b = torch.tensor([0.5])
    bias = torch.arange(6, dtype=torch.float32) if with_bias else None
    torch.library.opcheck(quack_fp8._quack_fp8_scaled_mm, (a, b, scale_a, scale_b, torch.bfloat16, bias))


def test_warmup_matches_inference_weight_layout(monkeypatch):
    calls = []

    def gemm(a, b, *, out, bias, alpha, tuned):
        assert torch.is_inference_mode_enabled()
        assert a.dtype == b.dtype == torch.float8_e4m3fn
        assert a.is_contiguous()
        assert b.stride() == (1, b.shape[0])
        assert bias is None
        assert tuned is True
        calls.append((a.shape[0], a.shape[1], b.shape[1]))
        out.zero_()

    monkeypatch.setattr(quack_fp8, "_gemm_interface", SimpleNamespace(gemm=gemm))
    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
    shapes = [(4, 8, 6), (2, 16, 8)]
    quack_fp8.warmup_quack_fp8(shapes, device="cpu")
    assert calls == shapes