# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Fixtures for the M0 local text mode.

Everything here is synthetic on purpose: a hardware profile built from
literals, and a checkpoint directory with a real ``config.json`` and
zero-filled weight files. The planner reads sizes and geometry, never tensors,
so this exercises the real code path without a GPU or 7.7 GB of weights. The
one test module that does need real weights says so and skips.
"""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from vllm_omni.edge.hardware_probe import (
    ACCEL_GPU_DISCRETE,
    ACCEL_GPU_INTEGRATED,
    ACCEL_NPU,
    HardwareProfile,
)

# A 28-layer Spark-shaped decoder: 7 full-attention and 21 sliding@512, which
# is the 1.7B layout. Small hidden/vocab so nothing here is slow.
SPARK_SHAPED_CONFIG = {
    "architectures": ["Spark2_5ForCausalLM"],
    "model_type": "spark2_5",
    "dtype": "bfloat16",
    "hidden_size": 2048,
    "intermediate_size": 8192,
    "num_hidden_layers": 28,
    "num_attention_heads": 8,
    "num_key_value_heads": 2,
    "head_dim": 256,
    "hidden_act": "gelu",
    "vocab_size": 131072,
    "sliding_window": 512,
    "max_position_embeddings": 1048576,
    "tie_word_embeddings": True,
    "layer_types": [
        "full_attention" if (i + 1) % 4 == 0 else "sliding_attention"
        for i in range(28)
    ],
    "rope_parameters": {
        "full_attention": {"rope_theta": 5000000, "partial_rotary_factor": 0.25},
        "sliding_attention": {"rope_theta": 10000, "partial_rotary_factor": 1.0},
    },
}

INT8_QUANT_CONFIG = {
    "quant_method": "compressed-tensors",
    "format": "int-quantized",
    "quantization_status": "compressed",
    "ignore": ["lm_head"],
    "config_groups": {
        "group_0": {
            "targets": ["Linear"],
            "weights": {"num_bits": 8, "type": "int", "symmetric": True, "strategy": "channel"},
        }
    },
}

FP8_QUANT_CONFIG = {
    "quant_method": "compressed-tensors",
    "format": "float-quantized",
    "config_groups": {},
}


def write_checkpoint(
    root: Path,
    *,
    quantization: dict | None = None,
    weight_bytes: int = 4 * 1024 * 1024,
    shards: int = 1,
) -> Path:
    """A checkpoint directory the manifest and planner can read end to end."""
    root.mkdir(parents=True, exist_ok=True)
    config = dict(SPARK_SHAPED_CONFIG)
    if quantization is not None:
        config["quantization_config"] = quantization
    (root / "config.json").write_text(json.dumps(config), encoding="utf-8")
    (root / "tokenizer_config.json").write_text('{"model_max_length": 4096}', encoding="utf-8")
    (root / "chat_template.jinja").write_text("{{ messages }}", encoding="utf-8")
    per_shard = max(weight_bytes // shards, 1)
    for i in range(shards):
        name = "model.safetensors" if shards == 1 else f"model-{i + 1:05d}-of-{shards:05d}.safetensors"
        with (root / name).open("wb") as fh:
            fh.truncate(per_shard)
    return root


@pytest.fixture
def dense_checkpoint(tmp_path: Path) -> Path:
    return write_checkpoint(tmp_path / "dense", weight_bytes=8 * 1024 * 1024, shards=2)


@pytest.fixture
def int8_checkpoint(tmp_path: Path) -> Path:
    return write_checkpoint(tmp_path / "int8", quantization=INT8_QUANT_CONFIG)


def blackwell_laptop(ram_available_bytes: int = 28 * 2**30) -> HardwareProfile:
    """The reference machine: sm_120 discrete GPU, an iGPU and an NPU beside it.

    Modelled on the probe output recorded in the architecture proposal, so the
    tests exercise the same shape the real host produces.
    """
    return HardwareProfile(
        arch="x86_64",
        cpu_model="AMD Ryzen AI 9 HX 370 w/ Radeon 890M",
        cpu_flags=["avx2", "avx512", "avx512_bf16", "avx512_vnni", "fma"],
        n_logical=24,
        n_physical=12,
        clusters=[{"max_khz": 5000000, "cpus": list(range(24))}],
        ram_total_bytes=int(30.91 * 2**30),
        ram_available_bytes=ram_available_bytes,
        accelerator="cuda_discrete",
        gpu_name="NVIDIA GeForce RTX 5090 Laptop GPU",
        gpu_mem_bytes=24463 * 2**20,
        gpu_sm_count=110,
        gpu_capability=[12, 0],
        supports_bf16=True,
        host_os="wsl2",
        accelerators=[
            {"kind": ACCEL_GPU_DISCRETE, "vendor": "nvidia",
             "name": "NVIDIA GeForce RTX 5090 Laptop GPU", "usable": True,
             "source": "torch.cuda", "mem_bytes": 24463 * 2**20},
            {"kind": ACCEL_NPU, "vendor": "amd", "name": "NPU Compute Accelerator Device",
             "usable": False, "source": "wsl-host", "reason": "no device node under WSL2"},
            {"kind": ACCEL_GPU_INTEGRATED, "vendor": "amd", "name": "AMD Radeon(TM) 890M Graphics",
             "usable": False, "source": "wsl-host", "reason": "no /dev/dri under WSL2"},
        ],
        source="test-fixture",
    )


def cpu_only_box() -> HardwareProfile:
    """The no-NVIDIA deployment class, as a machine rather than as a mask."""
    profile = blackwell_laptop()
    profile.accelerator = "none"
    profile.gpu_name = ""
    profile.gpu_mem_bytes = 0
    profile.gpu_capability = []
    profile.accelerators = [a for a in profile.accelerators if a["vendor"] != "nvidia"]
    return profile


@pytest.fixture
def laptop_profile() -> HardwareProfile:
    return blackwell_laptop()