"""Compare graphs generated by vLLM-Spyre vs AFTU

Run `python -m pytest tests/aftu/test_compare_graphs.py`.
"""

import os
import sys
import tempfile

import pytest
import torch
from graph_compare_utils import (collect_graph_files, compare_graphs,
                                 get_model_path,
                                 run_inference_py_and_get_graphs)
from pytest_mock.plugin import MockerFixture
from spyre_util import DecodeWarmupShapes, ModelInfo, patch_environment
from vllm import LLM

from vllm_spyre.model_executor.model_loader.spyre import SpyreCausalLM


# TODO: AFTU uses float32 for mask, while we use float16, which
# can generates different graphs. For now, until AFTU does not change
# this, we can just mock this method to match with AFTU behavior.
# NOTE: we need to set VLLM_ENABLE_V1_MULTIPROCESSING=0 otherwise this
# mock will not propagate to the child process of the model runner
def mock_get_mask_dtype(mocker: MockerFixture):

    mocker.patch.object(SpyreCausalLM,
                        "get_mask_dtype",
                        return_value=torch.float32)


@pytest.mark.spyre
@pytest.mark.cb
def test_compare_graphs_cb(model: ModelInfo, max_num_seqs: int,
                           max_model_len: int, monkeypatch: pytest.MonkeyPatch,
                           mocker: MockerFixture):
    """Test if the graphs generated by vllm-spyre matches with the ones by
    AFTU in continuous batching mode """

    model_path = get_model_path(model)

    attn_type = 'paged_fp8' if model.is_quantized \
        else 'paged'

    if model.is_quantized:
        mock_get_mask_dtype(mocker)

    inference_py_args = [
        sys.executable, "-m", "aiu_fms_testing_utils.scripts.inference",
        "--architecture", "hf_pretrained", "--model_path", model_path,
        "--tokenizer", model_path, "--unfuse_weights", "--device_type", "aiu",
        "--compile", "--cast_bf16_to_fp16", "--compile_dynamic",
        "--min_pad_length", "64", "--max_new_tokens", "5", "--batch_size",
        str(max_num_seqs), "--compile_dynamic_sendnn", "--attention_type",
        attn_type
    ]

    if not model.is_quantized:
        inference_py_args += ["--default_dtype", "fp16"]

    extra_env = {
        "VLLM_DT_MAX_CONTEXT_LEN": str(max_model_len),
        "VLLM_DT_MAX_BATCH_SIZE": str(max_num_seqs),
        "VLLM_DT_MAX_BATCH_TKV_LIMIT": str(1024 * 128)
    }
    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args, extra_env)

    assert len(aftu_graphs) > 0

    # VLLM

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_cb")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")

    # need for the mocker
    monkeypatch.setenv('VLLM_ENABLE_V1_MULTIPROCESSING', 0)
    patch_environment(use_cb=True,
                      warmup_shapes=None,
                      backend="sendnn",
                      monkeypatch=monkeypatch)

    original_cwd = os.getcwd()
    try:
        # Change to temp dir to set the test environment clean
        with tempfile.TemporaryDirectory() as tmpdir:
            os.chdir(tmpdir)

            # We only need to load the model
            LLM(model=model.name,
                revision=model.revision,
                max_model_len=max_model_len,
                tensor_parallel_size=1,
                max_num_seqs=max_num_seqs)

            vllm_graphs = collect_graph_files(tmpdir)
    finally:
        # Restore in case of exception
        os.chdir(original_cwd)

    assert compare_graphs(vllm_graphs, aftu_graphs)


@pytest.mark.spyre
@pytest.mark.parametrize(
    "warmup_shapes", [[(64, 4, 4)]])  # (prompt_length/new_tokens/batch_size)
def test_compare_graphs_static_batching(model: ModelInfo,
                                        warmup_shapes: DecodeWarmupShapes,
                                        monkeypatch: pytest.MonkeyPatch,
                                        mocker: MockerFixture) -> None:
    """Test if the graphs generated by vllm-spyre matches with the ones by
    AFTU in static batching mode """

    attn_type = 'math_fp8' if model.is_quantized \
        else 'sdpa'

    if model.is_quantized:
        mock_get_mask_dtype(mocker)

    model_path = get_model_path(model)

    inference_py_args = [
        sys.executable, "-m", "aiu_fms_testing_utils.scripts.inference",
        "--architecture", "hf_pretrained", "--model_path", model_path,
        "--tokenizer", model_path, "--unfuse_weights", "--device_type", "aiu",
        "--compile", "--cast_bf16_to_fp16", "--compile_dynamic",
        "--fixed_prompt_length",
        str(warmup_shapes[0][0]), "--max_new_tokens",
        str(warmup_shapes[0][1]), "--batch_size",
        str(warmup_shapes[0][2]), "--attention_type", attn_type
    ]

    if not model.is_quantized:
        inference_py_args += ["--default_dtype", "fp16"]

    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args)
    assert len(aftu_graphs) > 0
    # VLLM

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_sb")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")
    # needed for the mocker
    monkeypatch.setenv('VLLM_ENABLE_V1_MULTIPROCESSING', 0)
    patch_environment(use_cb=False,
                      warmup_shapes=warmup_shapes,
                      backend="sendnn",
                      monkeypatch=monkeypatch)

    original_cwd = os.getcwd()
    try:
        # Change to temp dir to set the test environment clean
        with tempfile.TemporaryDirectory() as tmpdir:
            os.chdir(tmpdir)

            LLM(model=model.name,
                revision=model.revision,
                max_model_len=2048,
                tensor_parallel_size=1,
                max_num_seqs=warmup_shapes[0][2])

            vllm_graphs = collect_graph_files(tmpdir)
    finally:
        # Restore in case of exception
        os.chdir(original_cwd)

    assert compare_graphs(vllm_graphs, aftu_graphs)


@pytest.mark.spyre
@pytest.mark.chunked_prefill
def test_compare_graphs_chunked_prefill(model: ModelInfo, max_num_seqs: int,
                                        max_model_len: int,
                                        monkeypatch: pytest.MonkeyPatch):
    """Test if the graphs generated by vllm-spyre matches with the ones by
    AFTU with chunked prefill enabled """

    if model.is_quantized:
        pytest.skip("Quantized model are not yet supported "
                    "with chunked prefill")
    model_path = get_model_path(model)

    chunk_size = 128
    inference_py_args = [
        sys.executable,
        "-m",
        "aiu_fms_testing_utils.scripts.inference",
        "--architecture",
        "hf_pretrained",
        "--model_path",
        model_path,
        "--tokenizer",
        model_path,
        "--unfuse_weights",
        "--device_type",
        "aiu",
        "--compile",
        "--cast_bf16_to_fp16",
        "--compile_dynamic",
        "--min_pad_length",
        "64",
        "--max_new_tokens",
        "5",
        "--batch_size",
        str(max_num_seqs),
        "--compile_dynamic_sendnn",
        "--attention_type",
        'paged',
        "--default_dtype",
        "fp16",
        "--prefill_chunk_size",
        str(chunk_size),
    ]

    extra_env = {
        "VLLM_DT_MAX_CONTEXT_LEN": str(max_model_len),
        "VLLM_DT_MAX_BATCH_SIZE": str(max_num_seqs),
        "VLLM_DT_MAX_BATCH_TKV_LIMIT": str(1024 * 128),
        "VLLM_DT_CHUNK_LEN": str(chunk_size)
    }
    aftu_graphs = run_inference_py_and_get_graphs(inference_py_args, extra_env)

    assert len(aftu_graphs) > 0

    # VLLM

    monkeypatch.setenv("DEE_DUMP_GRAPHS", "vllm_cb_cp")
    # Disable cache to produce the graphs
    monkeypatch.setenv("TORCH_SENDNN_CACHE_ENABLE", "0")

    monkeypatch.setenv('VLLM_DT_CHUNK_LEN', str(chunk_size))
    patch_environment(use_cb=True,
                      warmup_shapes=None,
                      backend="sendnn",
                      monkeypatch=monkeypatch,
                      use_chunked_prefill=True)

    original_cwd = os.getcwd()
    try:
        # Change to temp dir to set the test environment clean
        with tempfile.TemporaryDirectory() as tmpdir:
            os.chdir(tmpdir)

            # We only need to load the model
            LLM(model=model.name,
                revision=model.revision,
                max_model_len=max_model_len,
                tensor_parallel_size=1,
                max_num_batched_tokens=chunk_size,
                max_num_seqs=max_num_seqs)

            vllm_graphs = collect_graph_files(tmpdir)
    finally:
        # Restore in case of exception
        os.chdir(original_cwd)

    assert compare_graphs(vllm_graphs, aftu_graphs)
