from __future__ import annotations

from datetime import datetime, timezone
import json
import time

import pytest

from obliteratus.gpu_lifecycle import (
    GpuLifecyclePublisher,
    MemoryUsage,
    measure_torch_memory,
)


def test_fake_supervisor_observes_order_identity_and_recoverable_state(tmp_path):
    publisher = GpuLifecyclePublisher(
        tmp_path,
        heartbeat_seconds=0.01,
        clock=lambda: datetime(2026, 1, 2, tzinfo=timezone.utc),
        run_id="run-1",
    )
    memory = MemoryUsage(allocated_bytes=10, reserved_bytes=12, device_count=1)
    publisher.loading("org/model")
    publisher.resize(memory)
    publisher.ready(memory)
    time.sleep(1.05)
    publisher.release(reason="test_complete")
    assert publisher.release(reason="duplicate") is None

    events = [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()]
    assert [event["event"] for event in events] == [
        "loading", "resize", "ready", "heartbeat", "release",
    ]
    assert [event["sequence"] for event in events] == list(range(1, 6))
    assert len({event["event_id"] for event in events}) == 5
    assert all(event["run_id"] == "run-1" for event in events)
    assert events[1]["reserved_vram_bytes"] == 12
    current = json.loads((tmp_path / "current.json").read_text())
    assert current["event"] == "release"
    assert current["reason"] == "test_complete"


def test_disabled_publisher_is_noop():
    publisher = GpuLifecyclePublisher(None)
    assert publisher.loading("model") is None
    assert publisher.resize(MemoryUsage()) is None
    assert publisher.ready() is None
    assert publisher.heartbeat() is None
    assert publisher.release() is None


def test_runtime_directory_must_exist(tmp_path):
    with pytest.raises(ValueError, match="must already exist"):
        GpuLifecyclePublisher(tmp_path / "missing")


def test_runtime_directory_rejects_symlink(tmp_path):
    real = tmp_path / "real"
    real.mkdir()
    link = tmp_path / "link"
    link.symlink_to(real, target_is_directory=True)
    with pytest.raises(ValueError, match="must not be a symlink"):
        GpuLifecyclePublisher(link)


def test_measure_torch_memory_aggregates_devices():
    class FakeCuda:
        is_available = staticmethod(lambda: True)
        device_count = staticmethod(lambda: 2)
        memory_allocated = staticmethod(lambda index: (index + 1) * 10)
        memory_reserved = staticmethod(lambda index: (index + 1) * 20)

    usage = measure_torch_memory(type("Torch", (), {"cuda": FakeCuda})())
    assert usage == MemoryUsage(allocated_bytes=30, reserved_bytes=60, device_count=2)