import json
import os
import importlib.util
import sys
from pathlib import Path


def _load_jit_cache_module():
    repo_root = Path(__file__).resolve().parents[4]
    module_path = (
        repo_root
        / "python"
        / "sglang"
        / "srt"
        / "layers"
        / "deep_gemm_wrapper"
        / "jit_cache.py"
    )
    spec = importlib.util.spec_from_file_location("deep_gemm_jit_cache_under_test", module_path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def test_prepare_deep_gemm_jit_env_sets_cache_before_import(tmp_path, monkeypatch):
    prepare_deep_gemm_jit_env = _load_jit_cache_module().prepare_deep_gemm_jit_env

    cache_dir = tmp_path / "dg-cache"
    monkeypatch.delenv("DG_JIT_CACHE_DIR", raising=False)
    monkeypatch.delenv("DG_JIT_USE_NVRTC", raising=False)
    monkeypatch.setenv("SGLANG_DG_CACHE_DIR", str(cache_dir))
    monkeypatch.setenv("SGLANG_DG_USE_NVRTC", "1")

    resolved = prepare_deep_gemm_jit_env()

    assert resolved == str(cache_dir)
    assert os.environ["DG_JIT_CACHE_DIR"] == str(cache_dir)
    assert os.environ["DG_JIT_USE_NVRTC"] == "1"
    assert cache_dir.is_dir()


def test_prepare_deep_gemm_jit_env_respects_explicit_deepgemm_env(
    tmp_path, monkeypatch
):
    prepare_deep_gemm_jit_env = _load_jit_cache_module().prepare_deep_gemm_jit_env

    explicit = tmp_path / "explicit"
    ignored = tmp_path / "ignored"
    monkeypatch.setenv("DG_JIT_CACHE_DIR", str(explicit))
    monkeypatch.setenv("SGLANG_DG_CACHE_DIR", str(ignored))
    monkeypatch.setenv("DG_JIT_USE_NVRTC", "0")
    monkeypatch.setenv("SGLANG_DG_USE_NVRTC", "1")

    resolved = prepare_deep_gemm_jit_env()

    assert resolved == str(explicit)
    assert os.environ["DG_JIT_CACHE_DIR"] == str(explicit)
    assert os.environ["DG_JIT_USE_NVRTC"] == "0"
    assert explicit.is_dir()
    assert not ignored.exists()


def test_sparse_m_list_is_sparse_but_covers_boundaries():
    build_sparse_m_list = _load_jit_cache_module().build_sparse_m_list

    m_list = build_sparse_m_list(65536)

    assert m_list[0] == 1
    assert m_list[-1] == 65536
    assert 1024 in m_list
    assert 2048 in m_list
    assert 4096 in m_list
    assert len(m_list) < 5000
    assert len(m_list) < 65536 // 10


def test_build_warmup_m_lists_is_category_and_cp_aware():
    build_warmup_m_lists = _load_jit_cache_module().build_warmup_m_lists

    lists = build_warmup_m_lists(
        chunked_prefill_size=65536,
        attn_cp_size=8,
        max_running_requests=200,
        speculative_num_draft_tokens=4,
    )

    assert lists.per_rank[-1] == 8192
    assert lists.gathered[-1] == 65536
    assert lists.decode[-1] == 800
    assert len(lists.per_rank) < len(lists.gathered)
    assert len(lists.gathered) < 5000


def test_warmup_manifest_requires_cache_key_superset(tmp_path):
    jit_cache = _load_jit_cache_module()
    cache_entries = jit_cache.cache_entries
    make_warmup_manifest_key = jit_cache.make_warmup_manifest_key
    mark_warmup_complete = jit_cache.mark_warmup_complete
    should_skip_warmup = jit_cache.should_skip_warmup

    cache_root = tmp_path / "deep_gemm"
    (cache_root / "cache" / "a").mkdir(parents=True)
    (cache_root / "cache" / "b").mkdir()
    key = make_warmup_manifest_key(
        kernel_type="GEMM_NT_F8F8BF16",
        n=4096,
        k=8192,
        num_groups=1,
        m_list=[1, 2, 4, 8],
        runtime_fingerprint={"deep_gemm": "0.1.2", "sm": 90},
    )

    mark_warmup_complete(cache_root, key, cache_entries(cache_root))
    assert should_skip_warmup(cache_root, key)

    # Deleting a cubin directory must invalidate the SGLang-side manifest.
    (cache_root / "cache" / "b").rmdir()
    assert not should_skip_warmup(cache_root, key)


def test_warmup_manifest_file_is_json_and_atomic(tmp_path):
    jit_cache = _load_jit_cache_module()
    MANIFEST_FILENAME = jit_cache.MANIFEST_FILENAME
    mark_warmup_complete = jit_cache.mark_warmup_complete

    cache_root = tmp_path / "deep_gemm"
    cache_root.mkdir()
    mark_warmup_complete(cache_root, "shape-key", {"abc"})

    manifest = json.loads((cache_root / MANIFEST_FILENAME).read_text())
    assert manifest["version"] == 1
    assert manifest["entries"]["shape-key"]["status"] == "complete"
    assert manifest["entries"]["shape-key"]["cache_entries"] == ["abc"]