"""Unit tests for the pure calculation logic in vramcalc."""

import pytest

from vramcalc import (
    QUANT_BPW,
    weight_bytes,
    quant_from_filename,
    kv_cache_bytes,
    compute_scratch_bytes,
    yarn_effective_context,
    yarn_warnings,
    gpu_split,
    estimate,
    command_preview,
    format_bytes,
)
from vramcalc.presets import PRESETS
from vramcalc.gguf import parse_header_bytes, parse_header_with_tensors, metadata_to_arch


def test_quant_table_has_common_types():
    for q in ["Q4_K_M", "Q5_K_M", "Q8_0", "Q2_K", "F16", "BF16", "Q6_K"]:
        assert q in QUANT_BPW


def test_quant_table_has_rocmfp_types():
    for q in ["ROCmFP4", "ROCmFPX"]:
        assert q in QUANT_BPW


def test_quant_from_filename():
    assert quant_from_filename("Llama-3-8B-Instruct-Q4_K_M.gguf") == "Q4_K_M"
    assert quant_from_filename("model.Q5_K_S.gguf") == "Q5_K_S"
    assert quant_from_filename("model.Q8_0.gguf") == "Q8_0"
    assert quant_from_filename("model.F16.gguf") == "F16"
    assert quant_from_filename("model.gguf") is None


def test_quant_from_filename_rocmfp():
    # Real chadrock filenames — case-insensitive token match.
    assert quant_from_filename(
        "CHADROCK3.6-27B-Coder-MTP-ROCmFP4-STRIX_LEAN.gguf"
    ) == "ROCmFP4"
    assert quant_from_filename(
        "CHADROCK3.6-35B-A3B-Coder-MTP-ROCmFPX-MoEQuality-7.08BPW.gguf"
    ) == "ROCmFPX"
    # plain ROCm with no FP version must not falsely match ROCmFP4/FPX
    assert quant_from_filename("model-ROCM.gguf") is None


def test_weight_bytes_rocmfp4_matches_real_file():
    # 27B Strix Lean: 14,817,251,680 bytes for 27,320,697,856 tensor elems.
    # Effective bpw from the real file -> weight estimate reproduces file size.
    b = weight_bytes(27_320_697_856, "ROCmFP4")
    assert 14.6e9 < b < 15.0e9  # ~14.82 GiB file size


def test_weight_bytes():
    # 8B params at Q4_K_M (4.84375 bpw) -> ~4.84 GB
    b = weight_bytes(8_000_000_000, "Q4_K_M")
    assert 4.7e9 < b < 4.9e9
    # F16 -> 16 bytes/param
    assert weight_bytes(1_000_000, "F16") == pytest.approx(2_000_000)


def test_weight_bytes_unknown_raises():
    with pytest.raises(ValueError):
        weight_bytes(1, "BOGUS")


def test_kv_cache_basic():
    # Llama-3 8B: 32 layers, 4096 embd, 32 heads, 8 kv heads, ctx 8192, f16
    # head_dim = 128; per layer = 8192*2*8*128*2 = 33,554,432 bytes
    # * 32 layers = ~1.07 GB
    b = kv_cache_bytes(
        n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
        n_ctx=8192, cache_dtype="f16",
    )
    assert 1.0e9 < b < 1.1e9


def test_kv_cache_scales_with_ctx():
    b1 = kv_cache_bytes(n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
                        n_ctx=4096, cache_dtype="f16")
    b2 = kv_cache_bytes(n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
                        n_ctx=8192, cache_dtype="f16")
    assert b2 == pytest.approx(2 * b1)


def test_kv_cache_mtp():
    base = kv_cache_bytes(n_layer=61, n_embd=7168, n_head=128, n_head_kv=128,
                          n_ctx=8192, cache_dtype="f16", n_mtp=0)
    mtp = kv_cache_bytes(n_layer=61, n_embd=7168, n_head=128, n_head_kv=128,
                         n_ctx=8192, cache_dtype="f16", n_mtp=1)
    # n_mtp=1 adds one layer's worth: ratio = 62/61
    assert mtp == pytest.approx(base * 62 / 61)


def test_kv_cache_quantized_smaller():
    f16 = kv_cache_bytes(n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
                         n_ctx=8192, cache_dtype="f16")
    q8 = kv_cache_bytes(n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
                        n_ctx=8192, cache_dtype="q8_0")
    q4 = kv_cache_bytes(n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
                        n_ctx=8192, cache_dtype="q4_0")
    assert q8 < f16
    assert q4 < q8
    assert q8 == pytest.approx(f16 / 2)
    assert q4 == pytest.approx(f16 / 4)


def test_compute_scratch_positive():
    s = compute_scratch_bytes(
        n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
        n_batch=512, compute_dtype="f16",
    )
    assert s > 0


def test_compute_scratch_no_fa_quantized_adds_dequant():
    base = compute_scratch_bytes(
        n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
        n_batch=512, compute_dtype="f16", cache_dtype="q8_0", flash_attn=True,
    )
    nofa = compute_scratch_bytes(
        n_layer=32, n_embd=4096, n_head=32, n_head_kv=8,
        n_batch=512, compute_dtype="f16", cache_dtype="q8_0", flash_attn=False,
    )
    assert nofa > base


def test_yarn_effective_context():
    # scale 0.5 doubles effective context
    assert yarn_effective_context(8192, 0.5) == 16384
    assert yarn_effective_context(8192, 1.0) == 8192
    # zero/negative scale falls back to training ctx
    assert yarn_effective_context(8192, 0.0) == 8192


def test_yarn_warnings_target_exceeds():
    w = yarn_warnings(
        training_ctx=8192, target_ctx=32768,
        rope_freq_scale=1.0, yarn_ext_factor=-1.0, yarn_attn_factor=1.0,
    )
    assert any("exceeds training context" in x for x in w)
    assert any("Effective context" in x for x in w)


def test_yarn_warnings_no_warning_when_within():
    w = yarn_warnings(
        training_ctx=131072, target_ctx=8192,
        rope_freq_scale=1.0, yarn_ext_factor=-1.0, yarn_attn_factor=1.0,
    )
    assert w == []


def test_gpu_split_proportional_sums():
    res = gpu_split(
        gpu_vram_bytes=[24 << 30, 24 << 30, 16 << 30],
        weights_bytes=40 << 30,
        kv_compute_bytes=4 << 30,
    )
    # weight shares sum to weights
    total_w = sum(a.weight_bytes for a in res.assignments)
    assert total_w == pytest.approx(40 << 30, rel=1e-6)
    # KV on exactly one GPU (default GPU 0)
    kv_gpus = [a for a in res.assignments if a.kv_compute_bytes > 0]
    assert len(kv_gpus) == 1
    assert kv_gpus[0].index == 0
    # used+free == vram for each gpu
    for a in res.assignments:
        assert a.used_bytes + a.free_bytes == pytest.approx(a.vram_bytes, rel=1e-6)


def test_gpu_split_kv_on_largest():
    res = gpu_split(
        gpu_vram_bytes=[16 << 30, 24 << 30],
        weights_bytes=10 << 30,
        kv_compute_bytes=3 << 30,
        kv_on_largest=True,
    )
    kv_gpus = [a for a in res.assignments if a.kv_compute_bytes > 0]
    assert len(kv_gpus) == 1
    assert kv_gpus[0].index == 1


def test_gpu_split_fit_badge():
    res = gpu_split(
        gpu_vram_bytes=[4 << 30],  # tiny
        weights_bytes=10 << 30,
        kv_compute_bytes=1 << 30,
    )
    assert res.all_fit is False
    assert res.assignments[0].fits is False


def test_estimate_llama3_8b():
    arch = PRESETS["Llama-3 8B"]
    from vramcalc import Inputs
    bd = estimate(arch, Inputs(quant="Q4_K_M", n_ctx=8192, cache_dtype="f16",
                                flash_attn=True, compute_dtype="f16",
                                n_batch=512, gpu_vram_gb=[24.0]))
    # weights ~4.8 GB
    assert 4.5e9 < bd.weights_bytes < 5.0e9
    # KV ~1 GB
    assert 0.9e9 < bd.kv_cache_bytes < 1.1e9
    # total fits on 24 GB
    assert bd.gpu.all_fit is True
    assert bd.total_bytes < 24 << 30


def test_estimate_mtp_overhead():
    arch = PRESETS["DeepSeek-V3 (MoE)"]
    from vramcalc import Inputs
    bd = estimate(arch, Inputs(quant="Q4_K_M", n_ctx=8192, cache_dtype="f16",
                               gpu_vram_gb=[80.0]))
    assert bd.mtp_overhead_bytes > 0


def test_command_preview_basic():
    arch = PRESETS["Llama-3 8B"]
    from vramcalc import Inputs
    cmd = command_preview(arch, Inputs(quant="Q4_K_M", n_ctx=8192,
                                       cache_dtype="q8_0", flash_attn=True,
                                       gpu_vram_gb=[24.0]))
    assert "llama-server" in cmd
    assert "-c 8192" in cmd
    assert "--cache-type-k q8_0" in cmd
    assert "--flash-attn" in cmd


def test_command_preview_multigpu_split():
    arch = PRESETS["Llama-3 8B"]
    from vramcalc import Inputs
    cmd = command_preview(arch, Inputs(quant="Q4_K_M", n_ctx=8192,
                                       gpu_vram_gb=[24.0, 24.0]))
    assert "--tensor-split 24.0,24.0" in cmd


def test_command_preview_yarn_when_target_exceeds():
    arch = PRESETS["Llama-3 8B"]  # training_ctx 8192
    from vramcalc import Inputs
    cmd = command_preview(arch, Inputs(quant="Q4_K_M", n_ctx=32768,
                                       rope_freq_scale=0.5,
                                       gpu_vram_gb=[24.0]))
    assert "--rope-scaling yarn" in cmd
    assert "--rope-freq-scale 0.5" in cmd


def test_command_preview_rocmfp_uses_custom_runner():
    arch = PRESETS["Llama-3 8B"]
    from vramcalc import Inputs
    for q in ("ROCmFP4", "ROCmFPX"):
        cmd = command_preview(arch, Inputs(quant=q, n_ctx=8192,
                                           gpu_vram_gb=[24.0]))
        # custom runner, not stock llama-server
        assert cmd.startswith("rocmfpx-llama-server")
        assert "ROCmFPX" in cmd  # warning note present


def test_format_bytes():
    assert "GiB" in format_bytes(5 << 30)
    assert "MiB" in format_bytes(5 << 20)
    assert "KiB" in format_bytes(5 << 10)
    assert format_bytes(0).endswith("B")


def test_gguf_parse_header_minimal():
    # Build a tiny valid GGUF header in memory: magic, version, tensor_count=0,
    # kv_count=2, with one string and one int metadata entry.
    import struct
    def s(x): return x.encode("utf-8")
    buf = b""
    buf += struct.pack("<I", 0x46554747)  # GGUF
    buf += struct.pack("<I", 3)           # version
    buf += struct.pack("<Q", 0)           # tensor_count
    buf += struct.pack("<Q", 2)           # kv_count
    # kv 1: "general.architecture" = "llama"
    name = s("general.architecture")
    buf += struct.pack("<Q", len(name)) + name
    buf += struct.pack("<I", 8)  # STRING
    val = s("llama")
    buf += struct.pack("<Q", len(val)) + val
    # kv 2: "llama.block_count" = 32
    name2 = s("llama.block_count")
    buf += struct.pack("<Q", len(name2)) + name2
    buf += struct.pack("<I", 4)  # UINT32
    buf += struct.pack("<I", 32)
    meta = parse_header_bytes(buf)
    assert meta["general.architecture"] == "llama"
    arch = metadata_to_arch(meta)
    assert arch.architecture == "llama"
    assert arch.n_layer == 32
    # name defaults to "" when general.name is absent
    assert arch.name == ""


def test_gguf_parse_header_includes_general_name():
    import struct

    def s(x): return x.encode("utf-8")
    buf = b""
    buf += struct.pack("<I", 0x46554747)  # GGUF
    buf += struct.pack("<I", 3)           # version
    buf += struct.pack("<Q", 0)           # tensor_count
    buf += struct.pack("<Q", 2)           # kv_count
    # kv 1: "general.name" = "Test Model"
    k = s("general.name")
    buf += struct.pack("<Q", len(k)) + k
    buf += struct.pack("<I", 8)           # STRING
    v = s("Test Model")
    buf += struct.pack("<Q", len(v)) + v
    # kv 2: "general.architecture" = "qwen35"
    k2 = s("general.architecture")
    buf += struct.pack("<Q", len(k2)) + k2
    buf += struct.pack("<I", 8)           # STRING
    v2 = s("qwen35")
    buf += struct.pack("<Q", len(v2)) + v2
    meta = parse_header_bytes(buf)
    arch = metadata_to_arch(meta)
    assert arch.name == "Test Model"
    assert arch.architecture == "qwen35"


def test_gguf_params_fallback_from_tensor_elems():
    # No general.parameter_count; metadata_to_arch should fall back to summed
    # tensor element counts.
    import struct

    def s(x): return x.encode("utf-8")
    buf = b""
    buf += struct.pack("<I", 0x46554747)  # GGUF
    buf += struct.pack("<I", 3)           # version
    buf += struct.pack("<Q", 2)           # tensor_count = 2
    buf += struct.pack("<Q", 1)           # kv_count = 1
    # kv: "general.architecture" = "llama"
    k = s("general.architecture")
    buf += struct.pack("<Q", len(k)) + k
    buf += struct.pack("<I", 8)
    v = s("llama")
    buf += struct.pack("<Q", len(v)) + v
    # tensor 1: name "t1", 2 dims (4, 8) -> 32 elems
    n1 = s("t1")
    buf += struct.pack("<Q", len(n1)) + n1
    buf += struct.pack("<I", 2)           # n_dims
    buf += struct.pack("<Q", 4)
    buf += struct.pack("<Q", 8)
    buf += struct.pack("<I", 0)           # dtype F32
    buf += struct.pack("<Q", 0)           # offset
    # tensor 2: name "t2", 1 dim (68) -> 68 elems
    n2 = s("t2")
    buf += struct.pack("<Q", len(n2)) + n2
    buf += struct.pack("<I", 1)
    buf += struct.pack("<Q", 68)
    buf += struct.pack("<I", 0)
    buf += struct.pack("<Q", 32)
    meta, total = parse_header_with_tensors(buf)
    assert total == 32 + 68
    arch = metadata_to_arch(meta, total)
    # no general.parameter_count -> fallback kicks in
    assert arch.params == 100


def test_gguf_params_prefers_explicit_parameter_count():
    # When general.parameter_count is present, it wins over the tensor sum.
    import struct

    def s(x): return x.encode("utf-8")
    buf = b""
    buf += struct.pack("<I", 0x46554747)  # GGUF
    buf += struct.pack("<I", 3)           # version
    buf += struct.pack("<Q", 0)           # tensor_count
    buf += struct.pack("<Q", 1)           # kv_count
    k = s("general.parameter_count")
    buf += struct.pack("<Q", len(k)) + k
    buf += struct.pack("<I", 10)          # UINT64
    buf += struct.pack("<Q", 7_000_000_000)
    meta, total = parse_header_with_tensors(buf)
    assert meta["general.parameter_count"] == 7_000_000_000
    arch = metadata_to_arch(meta, total_tensor_elems=123)  # fallback ignored
    assert arch.params == 7_000_000_000


def test_gguf_tensor_section_unreachable_returns_minus_one():
    # Buffer too small to reach the tensor section -> total == -1, no fallback.
    import struct

    def s(x): return x.encode("utf-8")
    buf = b""
    buf += struct.pack("<I", 0x46554747)
    buf += struct.pack("<I", 3)
    buf += struct.pack("<Q", 3)           # tensor_count > 0
    buf += struct.pack("<Q", 1)           # kv_count
    k = s("general.architecture")
    buf += struct.pack("<Q", len(k)) + k
    buf += struct.pack("<I", 8)
    v = s("llama")
    buf += struct.pack("<Q", len(v)) + v
    # truncate so the tensor section is unreachable (don't append tensors)
    meta, total = parse_header_with_tensors(buf)
    assert total == -1
    arch = metadata_to_arch(meta, total)
    assert arch.params == 0  # negative fallback is treated as no fallback