import json
import time
from collections.abc import Iterator
from pathlib import Path

import pytest
import torch.distributed as dist

from deepinfer import envs
from deepinfer.config import ParallelConfig
from deepinfer.worker.decode_perf_cache import DecodePerfCache

pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]


@pytest.fixture
def cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[DecodePerfCache]:
    monkeypatch.setattr(envs, "DEEPINFER_DECODE_PERF_CACHE_DIR", str(tmp_path / "cache"))
    dist.init_process_group("gloo", init_method=(tmp_path / "rendezvous").as_uri(), rank=0, world_size=1)
    try:
        config = ParallelConfig(1, 1, 1, 1, 0, 0)
        yield DecodePerfCache("/models/model", config, [1], dist.group.WORLD)
    finally:
        dist.destroy_process_group()


@pytest.mark.parametrize("batch_size", ["1e309", "-1e309", "Infinity", "-Infinity"])
@pytest.mark.parametrize("has_older_curve", [False, True])
def test_overflowing_batch_size_is_a_cache_miss(cache: DecodePerfCache, batch_size: str, has_older_curve: bool) -> None:
    folder = Path(cache.cache_dir)
    folder.mkdir()
    now = int(time.time())
    if has_older_curve:
        (folder / f"{cache.hash_key}.{'a' * 32}.{now - 1}").write_text(json.dumps([[1, 2.5]]))
    (folder / f"{cache.hash_key}.{'b' * 32}.{now}").write_text(f"[[{batch_size}, 3.0]]")
    assert cache.load_decode_perf_cache() == ([(1, 2.5)] if has_older_curve else None)