"""
Tests for GateWay checkpointing system.
"""

import tempfile
from datetime import datetime
from pathlib import Path

import pytest


class TestCheckpointMetadata:
    """Tests for CheckpointMetadata."""

    def test_checkpoint_metadata_creation(self):
        """Test CheckpointMetadata creation."""
        from verlex_server.gateway.core.checkpointing import CheckpointMetadata, Framework

        metadata = CheckpointMetadata(
            checkpoint_id="ckpt_abc",
            job_id="job_123",
            timestamp=datetime.utcnow(),
            framework=Framework.PYTORCH,
            epoch=5,
            step=1000,
        )

        assert metadata.job_id == "job_123"
        assert metadata.checkpoint_id == "ckpt_abc"
        assert metadata.framework == Framework.PYTORCH
        assert metadata.epoch == 5
        assert metadata.step == 1000

    def test_checkpoint_metadata_to_dict(self):
        """Test CheckpointMetadata to_dict conversion."""
        from verlex_server.gateway.core.checkpointing import CheckpointMetadata, Framework

        metadata = CheckpointMetadata(
            checkpoint_id="ckpt_123",
            job_id="job_456",
            timestamp=datetime.utcnow(),
            framework=Framework.GENERIC,
        )

        d = metadata.to_dict()
        assert d["checkpoint_id"] == "ckpt_123"
        assert d["job_id"] == "job_456"


class TestFramework:
    """Tests for Framework enum."""

    def test_framework_values(self):
        """Test Framework enum values."""
        from verlex_server.gateway.core.checkpointing import Framework

        assert Framework.PYTORCH.value == "pytorch"
        assert Framework.TENSORFLOW.value == "tensorflow"
        assert Framework.JAX.value == "jax"
        assert Framework.GENERIC.value == "generic"


class TestLocalStorageBackend:
    """Tests for LocalStorageBackend."""

    def test_local_storage_init(self):
        """Test LocalStorageBackend initialization."""
        from verlex_server.gateway.core.checkpointing import LocalStorageBackend

        with tempfile.TemporaryDirectory() as tmpdir:
            backend = LocalStorageBackend(tmpdir)
            assert backend.base_path.exists()

    def test_local_storage_upload_download(self):
        """Test upload and download."""
        from verlex_server.gateway.core.checkpointing import LocalStorageBackend

        with tempfile.TemporaryDirectory() as tmpdir:
            backend = LocalStorageBackend(tmpdir)

            # Create a test file
            test_file = Path(tmpdir) / "test_source.bin"
            test_data = b"test checkpoint data"
            test_file.write_bytes(test_data)

            # Upload
            backend.upload(str(test_file), "test_ckpt/data.bin")

            # Download to different location
            download_path = Path(tmpdir) / "downloaded.bin"
            backend.download("test_ckpt/data.bin", str(download_path))

            assert download_path.read_bytes() == test_data

    def test_local_storage_delete(self):
        """Test file deletion."""
        from verlex_server.gateway.core.checkpointing import LocalStorageBackend

        with tempfile.TemporaryDirectory() as tmpdir:
            backend = LocalStorageBackend(tmpdir)

            # Create a test file
            test_file = Path(tmpdir) / "test_source.bin"
            test_file.write_bytes(b"data")

            # Upload
            backend.upload(str(test_file), "to_delete.bin")

            # Verify it exists
            assert (backend.base_path / "to_delete.bin").exists()

            # Delete
            backend.delete("to_delete.bin")

            # Verify deleted
            assert not (backend.base_path / "to_delete.bin").exists()


class TestGenericAdapter:
    """Tests for GenericAdapter."""

    def test_generic_adapter_save_load(self):
        """Test GenericAdapter save and load with file path."""
        from verlex_server.gateway.core.checkpointing import GenericAdapter

        adapter = GenericAdapter()

        state = {
            "model_weights": [1.0, 2.0, 3.0],
            "optimizer_state": {"lr": 0.001},
            "epoch": 10,
        }

        with tempfile.TemporaryDirectory() as tmpdir:
            ckpt_path = Path(tmpdir) / "checkpoint.pkl"

            # Save
            adapter.save(state, str(ckpt_path))
            assert ckpt_path.exists()

            # Load
            restored = adapter.load(str(ckpt_path))
            assert restored["epoch"] == 10
            assert restored["model_weights"] == [1.0, 2.0, 3.0]

    def test_generic_adapter_framework_property(self):
        """Test GenericAdapter framework property."""
        from verlex_server.gateway.core.checkpointing import Framework, GenericAdapter

        adapter = GenericAdapter()
        assert adapter.framework == Framework.GENERIC


class TestCheckpointer:
    """Tests for main Checkpointer class."""

    def test_checkpointer_initialization(self):
        """Test Checkpointer initialization."""
        from verlex_server.gateway.core.checkpointing import (
            Checkpointer,
            Framework,
            LocalStorageBackend,
        )

        with tempfile.TemporaryDirectory() as tmpdir:
            backend = LocalStorageBackend(tmpdir)

            checkpointer = Checkpointer(
                job_id="job_test",
                framework=Framework.GENERIC,
                storage_backend=backend,
            )

            assert checkpointer.job_id == "job_test"
            assert checkpointer.framework == Framework.GENERIC

    def test_checkpointer_with_custom_functions(self):
        """Test Checkpointer with custom save/restore functions."""
        from verlex_server.gateway.core.checkpointing import (
            Checkpointer,
            Framework,
            LocalStorageBackend,
        )

        with tempfile.TemporaryDirectory() as tmpdir:
            backend = LocalStorageBackend(tmpdir)

            saved_state = {"value": 42}

            def custom_save():
                return saved_state

            def custom_restore(state):
                saved_state.update(state)

            checkpointer = Checkpointer(
                job_id="job_custom",
                framework=Framework.GENERIC,
                storage_backend=backend,
                custom_save_fn=custom_save,
                custom_restore_fn=custom_restore,
            )

            assert checkpointer.custom_save_fn is not None
            assert checkpointer.custom_restore_fn is not None


class TestFrameworkDetection:
    """Tests for framework detection."""

    def test_detect_framework_function_exists(self):
        """Test detect_framework function exists."""
        from verlex_server.gateway.core.checkpointing import detect_framework

        assert callable(detect_framework)

    def test_detect_framework_default_generic(self):
        """Test default framework is generic."""
        from verlex_server.gateway.core.checkpointing import Framework, detect_framework

        def simple_func():
            return 42

        # Without actual ML imports, should default to generic
        result = detect_framework(simple_func)
        assert result in [Framework.GENERIC, Framework.PYTORCH, Framework.TENSORFLOW, Framework.JAX]


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