"""
Tests for Memory Manager utilities.
"""

import pytest
import torch

from dokodemo_ai.utils.memory import MemoryManager, MemoryState


class TestMemoryState:
    def test_gpu_utilization_zero_total(self):
        state = MemoryState(gpu_total_bytes=0, gpu_used_bytes=0)
        assert state.gpu_utilization == 0.0

    def test_gpu_utilization_half(self):
        state = MemoryState(
            gpu_total_bytes=1024**3,
            gpu_used_bytes=512 * 1024**2,
        )
        assert abs(state.gpu_utilization - 0.5) < 1e-6

    def test_summary_no_crash(self):
        state = MemoryState(
            gpu_total_bytes=4 * 1024**3,
            gpu_used_bytes=1 * 1024**3,
            gpu_free_bytes=3 * 1024**3,
            cpu_total_bytes=16 * 1024**3,
            cpu_available_bytes=8 * 1024**3,
            cpu_used_bytes=8 * 1024**3,
        )
        summary = state.summary()
        assert "GPU" in summary or "CPU" in summary


class TestMemoryManager:
    def test_device_auto_detect(self):
        mgr = MemoryManager()
        assert mgr.device is not None
        assert isinstance(mgr.device, torch.device)

    def test_explicit_cpu_device(self):
        mgr = MemoryManager(device=torch.device("cpu"))
        assert mgr.device.type == "cpu"

    def test_weight_budget_positive(self):
        mgr = MemoryManager(device=torch.device("cpu"))
        budget = mgr.get_weight_budget_bytes()
        assert budget > 0

    def test_kv_cache_estimate(self):
        mgr = MemoryManager()
        # 32 layers, 32 heads, 128 head_dim, 512 seq_len
        size = mgr.estimate_kv_cache_bytes(
            num_layers=32,
            hidden_size=4096,
            num_heads=32,
            max_seq_len=512,
        )
        # Should be >0 and <1GB for these params
        assert 0 < size < 1 * 1024**3

    def test_activation_estimate(self):
        mgr = MemoryManager()
        size = mgr.estimate_activation_bytes(
            hidden_size=4096,
            batch_size=1,
            seq_len=100,
        )
        assert size > 0

    def test_feasibility_cpu_mmap_always_true(self):
        """In mmap mode, any model size is feasible as long as ~1 GB active RAM exists."""
        mgr = MemoryManager(device=torch.device("cpu"))
        feasible, msg = mgr.check_feasibility(
            max_layer_size_bytes=50 * 1024**3,  # 50 GB model — fine in mmap mode
            embedding_size_bytes=1 * 1024**3,
            head_size_bytes=1 * 1024**3,
            use_mmap=True,
        )
        # mmap mode: OS page cache handles weights; only ~1 GB active RAM needed
        assert feasible
        assert "mmap" in msg.lower() or "page cache" in msg.lower() or "feasible" in msg.lower()

    def test_feasibility_cpu_small_layer_true(self):
        """Non-mmap CPU mode: a tiny layer easily fits in available RAM."""
        mgr = MemoryManager(device=torch.device("cpu"))
        feasible, msg = mgr.check_feasibility(
            max_layer_size_bytes=1 * 1024**2,  # 1 MB — trivially fits
            embedding_size_bytes=1 * 1024**2,
            head_size_bytes=1 * 1024**2,
            use_mmap=False,
        )
        assert feasible

    def test_cleanup_gpu_no_crash(self):
        mgr = MemoryManager(device=torch.device("cpu"))
        freed = mgr.cleanup_gpu()
        assert freed == 0  # CPU has nothing to free

    def test_parse_memory_string(self):
        """_parse_memory_string is tested indirectly via AutoModel"""
        from dokodemo_ai.auto_model import _parse_memory_string
        assert _parse_memory_string("4GB") == 4 * 1024**3
        assert _parse_memory_string("512MB") == 512 * 1024**2
        assert _parse_memory_string("2.5GB") == int(2.5 * 1024**3)
        assert _parse_memory_string("1TB") == 1024**4
