from app.memory import GB, estimate, fit_class, kv_cache_bytes, usable_gpu_bytes, weights_bytes
from app.parsing import detect_arch

from .test_parsing import QWEN3_8B_CONFIG


def test_kv_cache_qwen3_8b_32k_by_hand():
    # 2 (K,V) * 36 layers * 8 kv heads * 128 head_dim * 32768 tokens * 2 bytes
    expected = 2 * 36 * 8 * 128 * 32768 * 2
    kv, upper = kv_cache_bytes(detect_arch(QWEN3_8B_CONFIG), 32768)
    assert kv == expected and not upper
    assert round(kv / GB, 2) == 4.5


def test_sliding_window_caps_kv():
    cfg = {"num_hidden_layers": 4, "num_attention_heads": 8, "num_key_value_heads": 8, "head_dim": 64,
           "sliding_window": 128, "layer_types": ["sliding_attention", "full_attention"] * 2}
    arch = detect_arch(cfg)
    per = 2 * 8 * 64 * 2
    kv, _ = kv_cache_bytes(arch, 32768)
    assert kv == 2 * 32768 * per + 2 * 128 * per


def test_missing_config_is_flagged_rough():
    est = estimate(params=8e9, bits=4, mode="affine", file_bytes=None, arch=detect_arch(None),
                   context=32768, ram_gb=16)
    assert est.kv_rough and est.kv_gb and est.weights_source == "params"


def test_weights_prefer_files():
    assert weights_bytes(4_619_257_828, 8e9, 4) == (4_619_257_828, "files")
    b, src = weights_bytes(None, 8e9, 4)
    assert src == "params" and abs(b / 1e9 - 4.5) < 0.01
    assert weights_bytes(None, None, 4) == (None, "unknown")


def test_usable_memory_and_fit():
    assert usable_gpu_bytes(16) < usable_gpu_bytes(24) < usable_gpu_bytes(64)
    assert fit_class(int(5 * GB), 36) == "Comfortable"
    assert fit_class(int(12 * GB), 16) == "Unlikely"
    assert fit_class(None, 16) is None and fit_class(int(GB), None) is None


def test_context_changes_estimate_and_model_limit():
    arch = detect_arch(QWEN3_8B_CONFIG)
    small = estimate(params=8e9, bits=4, mode=None, file_bytes=None, arch=arch, context=4096, ram_gb=16)
    big = estimate(params=8e9, bits=4, mode=None, file_bytes=None, arch=arch, context=131072, ram_gb=16)
    assert big.total_gb > small.total_gb
    assert small.fit in ("Comfortable", "Likely") and big.fit == "Unlikely"
    assert big.exceeds_model_context and not small.exceeds_model_context