# Test Plan: Disk Cache Verification using PostProc Callback
#
# Purpose: Reliably test disk cache in CI by using postproc callbacks to detect
#          whether compilation actually happened or cache was used.
#
# Strategy:
# - postproc is ONLY called during codegen (cache miss)
# - postproc is NOT called when loading from cache (cache hit)
# - Use a counter in postproc to distinguish these cases
#
# CI Safety:
# - Use isolated cache/tmp directories per test (pytest tmp_path)
# - Use unique kernel identifiers (UUID + global_symbol) to avoid collisions
# - Clear memory cache between passes to force disk I/O
# - os.replace() requires source and dest on same filesystem (atomic rename)
#
# Technical Details:
# - Cache key is based on func.script(show_meta=True) hash
# - Python comments do NOT affect cache key (not in TIR)
# - Must use .with_attr("global_symbol", ...) to create unique cache keys

import pytest
import tilelang
import tilelang.testing
import tilelang.language as T
import tvm_ffi
import torch
import uuid
from pathlib import Path
from tilelang.env import env
from tilelang.cache import _dispatch_map

BACKENDS = [
    "tvm_ffi",
    "cython",
    "nvrtc",
    "cutedsl",
]


def _get_target_from_backend(backend: str):
    """Map backend to target string."""
    return "cutedsl" if backend == "cutedsl" else "auto"


def _require_backend_available(backend: str) -> None:
    if backend != "cutedsl":
        return
    try:
        from tilelang.jit.adapter.cutedsl.checks import check_cutedsl_available

        check_cutedsl_available()
    except ImportError as e:
        pytest.skip(f"CuTeDSL backend unavailable: {e}")


class PostProcCounter:
    """Track postproc callback invocations with a simple counter."""

    def __init__(self):
        self.count = 0
        self.marker = None

    def register_callback(self, backend: str):
        """Register postproc callback for the given backend."""
        comment_prefix = "#" if backend == "cutedsl" else "//"
        global_func = "tilelang_callback_cutedsl_postproc" if backend == "cutedsl" else "tilelang_callback_cuda_postproc"

        def callback(code, _):
            self.count += 1
            self.marker = f"{comment_prefix} CACHE_TEST_MARKER_{self.count}"
            return self.marker + "\n" + code

        tvm_ffi.register_global_func(global_func, f=callback, override=True)
        return callback


@pytest.fixture(scope="module", autouse=True)
def setup_module_env():
    """Setup and restore module-level environment and cache state."""
    # Save original env values
    original_cache_dir = env.TILELANG_CACHE_DIR

    # Enable cache once for entire module
    tilelang.enable_cache()

    yield

    # Restore env at module end
    env.TILELANG_CACHE_DIR = original_cache_dir

    # Restore default postproc callbacks
    tvm_ffi.register_global_func("tilelang_callback_cuda_postproc", f=lambda code, _: code, override=True)
    tvm_ffi.register_global_func("tilelang_callback_cutedsl_postproc", f=lambda code, _: code, override=True)


@pytest.fixture(scope="function")
def clean_cache_env(tmp_path, request):
    """Provide isolated cache environment for each test.

    Creates an isolated cache directory to ensure:
    - No interference from previous test runs
    - No interference between parallel tests
    - Clean slate for testing cache miss/hit behavior
    - Env restoration is handled by setup_module_env at module scope
    """
    # This fixture should ONLY be used with @pytest.mark.parametrize("backend", ...)
    backend = request.node.callspec.params["backend"]  # Will raise KeyError if missing
    _require_backend_available(backend)

    cache_dir = tmp_path / "tilelang_cache"
    cache_dir.mkdir()

    # Patch the cache root to point to an isolated directory.
    env.TILELANG_CACHE_DIR = str(cache_dir)

    # Clear memory caches to force disk I/O
    _dispatch_map[backend]._memory_cache.clear()

    return cache_dir


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("backend", BACKENDS)
def test_disk_cache_with_postproc(clean_cache_env, backend):
    """Test disk cache for multiple backends using postproc callback.

    Tests all CUDA-based backends: nvrtc, cutedsl
    (tvm_ffi, cython, torch use the same cuda_postproc callback as nvrtc)

    Verification logic:
    1. Pass 1: cache miss → postproc called → marker in kernel source
    2. Pass 2: cache hit → postproc NOT called → same marker still in source
    3. Verify cache files created on disk
    4. Verify functional correctness
    """
    counter = PostProcCounter()
    counter.register_callback(backend)

    # Use UUID in global_symbol to ensure unique cache key per test run
    unique_id = uuid.uuid4().hex[:8]
    M, N = 256, 256

    @T.prim_func
    def vector_add(
        A: T.Tensor((M, N), T.float32),
        B: T.Tensor((M, N), T.float32),
        C: T.Tensor((M, N), T.float32),
    ):
        with T.Kernel(M, threads=256) as bx:
            for i in T.serial(N):
                C[bx, i] = A[bx, i] + B[bx, i]

    kernel_func = vector_add.with_attr("global_symbol", f"vector_add_{backend}_{unique_id}")

    # === Pass 1: Cache miss (memory cache already cleared by fixture) ===
    kernel1 = tilelang.compile(
        kernel_func,
        out_idx=[2],
        target=_get_target_from_backend(backend),
        execution_backend=backend,
    )

    assert counter.count == 1, f"Cache miss: postproc should be called once, got {counter.count}"

    source1 = kernel1.get_kernel_source()
    assert counter.marker in source1, f"Expected marker '{counter.marker}' in kernel source"

    # Verify cache files created
    cache_files = list(Path(clean_cache_env).rglob("*.*"))
    assert len(cache_files) > 0, "Cache files should be created, found none"

    # === Pass 2: Cache hit (clear memory cache to force disk read) ===
    _dispatch_map[backend]._memory_cache.clear()

    kernel2 = tilelang.compile(
        kernel_func,
        out_idx=[2],
        target=_get_target_from_backend(backend),
        execution_backend=backend,
    )

    assert counter.count == 1, f"Cache hit: postproc should not be called again, got {counter.count} calls"
    if backend == "nvrtc":
        assert kernel1.adapter.kernels is not kernel2.adapter.kernels

    source2 = kernel2.get_kernel_source()
    assert counter.marker in source2, f"Expected cached marker '{counter.marker}' in source"

    # === Verify functional correctness ===
    a = torch.randn(M, N, dtype=torch.float32).cuda()
    b = torch.randn(M, N, dtype=torch.float32).cuda()

    c1 = kernel1(a, b)
    c2 = kernel2(a, b)
    ref = a + b

    torch.testing.assert_close(c1, ref)
    torch.testing.assert_close(c2, ref)
    torch.testing.assert_close(c1, c2)


@tilelang.testing.requires_cuda
def test_tvm_ffi_export_library_after_disk_cache_hit(tmp_path, monkeypatch):
    cache_dir = tmp_path / "tilelang_cache"
    cache_dir.mkdir()
    monkeypatch.setattr(env, "TILELANG_CACHE_DIR", str(cache_dir))

    cache = _dispatch_map["tvm_ffi"]
    cache._memory_cache.clear()

    @T.prim_func
    def vector_add(
        A: T.Tensor((128,), T.float32),
        B: T.Tensor((128,), T.float32),
    ):
        with T.Kernel(128, threads=128) as i:
            B[i] = A[i] + 1.0

    kernel_func = vector_add.with_attr("global_symbol", f"export_library_{uuid.uuid4().hex[:8]}")
    cold_kernel = tilelang.compile(kernel_func, out_idx=[1], execution_backend="tvm_ffi")
    cold_library = tmp_path / "cold.so"
    cold_kernel.export_library(str(cold_library))

    cache._memory_cache.clear()
    cached_kernel = tilelang.compile(kernel_func, out_idx=[1], execution_backend="tvm_ffi")
    cached_library = tmp_path / "cached.so"
    cached_kernel.export_library(str(cached_library))

    assert cold_library.is_file()
    assert cached_kernel.artifact is None
    assert cached_library.read_bytes() == Path(cached_kernel.adapter.libpath).read_bytes()
    assert tilelang.tvm.runtime.load_module(str(cached_library)) is not None


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("backend", BACKENDS)
def test_cache_miss_detection(clean_cache_env, backend):
    """Verify cache correctly misses when function changes.

    This ensures our testing method is valid - different functions should
    produce different cache keys and trigger recompilation.
    """
    counter = PostProcCounter()
    counter.register_callback(backend)

    M, N = 128, 128

    # Kernel 1: A + 1.0
    @T.prim_func
    def func1(A: T.Tensor((M, N), T.float32), B: T.Tensor((M, N), T.float32)):
        with T.Kernel(M, threads=128) as bx:
            for i in T.serial(N):
                B[bx, i] = A[bx, i] + 1.0

    unique_id_1 = uuid.uuid4().hex[:8]
    kernel_func1 = func1.with_attr("global_symbol", f"func1_{backend}_{unique_id_1}")

    tilelang.compile(
        kernel_func1,
        out_idx=[1],
        target=_get_target_from_backend(backend),
        execution_backend=backend,
    )
    assert counter.count == 1, f"First kernel: expected 1 call, got {counter.count}"

    # Kernel 2: A + 2.0 (different implementation)
    @T.prim_func
    def func2(A: T.Tensor((M, N), T.float32), B: T.Tensor((M, N), T.float32)):
        with T.Kernel(M, threads=128) as bx:
            for i in T.serial(N):
                B[bx, i] = A[bx, i] + 2.0  # Different!

    unique_id_2 = uuid.uuid4().hex[:8]
    kernel_func2 = func2.with_attr("global_symbol", f"func2_{backend}_{unique_id_2}")

    tilelang.compile(
        kernel_func2,
        out_idx=[1],
        target=_get_target_from_backend(backend),
        execution_backend=backend,
    )

    assert counter.count == 2, f"Different function should cause cache miss, expected 2 calls, got {counter.count}"


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("backend", BACKENDS)
def test_cache_isolation_between_tests(clean_cache_env, backend):
    """Verify cache isolation between tests.

    Ensures clean_cache_env fixture provides independent cache directory for each test.
    """
    # Verify cache directory is empty
    cache_files = list(Path(clean_cache_env).rglob("*"))
    assert all(f.is_dir() for f in cache_files), f"Cache should be empty, found: {cache_files}"

    # Compile a kernel
    counter = PostProcCounter()
    counter.register_callback(backend)

    unique_id = uuid.uuid4().hex[:8]

    @T.prim_func
    def simple(A: T.Tensor((128,), T.float32), B: T.Tensor((128,), T.float32)):
        with T.Kernel(128, threads=128) as i:
            B[i] = A[i] * 2.0

    kernel_func = simple.with_attr("global_symbol", f"simple_{backend}_{unique_id}")

    tilelang.compile(
        kernel_func,
        out_idx=[1],
        target=_get_target_from_backend(backend),
        execution_backend=backend,
    )

    # Should be cache miss (empty cache dir)
    assert counter.count == 1, f"Expected cache miss, got count={counter.count}"

    # Verify cache files created
    cache_files_after = list(Path(clean_cache_env).rglob("*.*"))
    assert len(cache_files_after) > 0, f"Cache files should be created, found: {cache_files_after}"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])