import pytest

from heapr.model_utils import build_max_memory, discover_sparse_layers, validate_model_device_placement


class TensorLike:
    def __init__(self, shape):
        self.shape = shape


class Experts:
    def __init__(self):
        self.gate_up_proj = TensorLike((256, 1024, 2048))
        self.down_proj = TensorLike((256, 2048, 512))


class SparseMlp:
    def __init__(self):
        self.experts = Experts()


class DenseMlp:
    pass


class Layer:
    def __init__(self, mlp):
        self.mlp = mlp


class Inner:
    def __init__(self):
        self.layers = [Layer(DenseMlp()), Layer(SparseMlp())]


class Model:
    def __init__(self):
        self.model = Inner()

    def named_modules(self):
        yield "", self
        yield "model.layers.0.mlp", self.model.layers[0].mlp
        yield "model.layers.1.mlp", self.model.layers[1].mlp


def test_discover_sparse_layers_from_packed_laguna_tensors():
    info = discover_sparse_layers(Model())

    assert len(info) == 1
    assert info[0].model_layer_idx == 1
    assert info[0].sparse_idx == 0
    assert info[0].num_experts == 256
    assert info[0].routed_width == 512
    assert info[0].hidden_size == 2048


class FakeCuda:
    def __init__(self, count=4):
        self._count = count

    def is_available(self):
        return True

    def device_count(self):
        return self._count


class FakeTorch:
    cuda = FakeCuda()


def test_build_max_memory_uses_all_visible_gpus_without_cpu(monkeypatch):
    monkeypatch.setattr("heapr.model_utils.require_torch", lambda: FakeTorch)

    max_memory = build_max_memory(gpu_memory_per_device="46GiB")

    assert max_memory == {0: "46GiB", 1: "46GiB", 2: "46GiB", 3: "46GiB"}


def test_build_max_memory_adds_cpu_only_when_allowed(monkeypatch):
    monkeypatch.setattr("heapr.model_utils.require_torch", lambda: FakeTorch)

    max_memory = build_max_memory(
        gpu_memory_per_device="46GiB",
        max_cpu_memory="80GiB",
        allow_cpu_offload=True,
    )

    assert max_memory["cpu"] == "80GiB"


def test_build_max_memory_rejects_cpu_memory_without_offload():
    with pytest.raises(ValueError, match="--allow-cpu-offload"):
        build_max_memory(gpu_memory_per_device="46GiB", max_cpu_memory="80GiB")


class FakeParameter:
    def __init__(self, device):
        self.device = device


class FakePlacedModel:
    def __init__(self, devices, hf_device_map=None):
        self.hf_device_map = hf_device_map
        self._parameters = [FakeParameter(device) for device in devices]

    def parameters(self):
        return iter(self._parameters)


def test_validate_model_device_placement_rejects_cpu_offload():
    model = FakePlacedModel(["cuda:0"], {"model.layers.0": 0, "model.layers.1": "cpu"})

    with pytest.raises(RuntimeError, match="offloaded"):
        validate_model_device_placement(model, allow_cpu_offload=False)


def test_validate_model_device_placement_rejects_single_gpu_when_multi_requested():
    model = FakePlacedModel(["cuda:0", "cuda:0"], {"": 0})

    with pytest.raises(RuntimeError, match="multi-GPU"):
        validate_model_device_placement(model, requested_gpu_count=4)


def test_validate_model_device_placement_accepts_multi_gpu_cuda_only():
    model = FakePlacedModel(["cuda:0", "cuda:1"], {"model.embed": 0, "model.layers.1": 1})

    validate_model_device_placement(model, requested_gpu_count=2)