"""
Tests for GateWay Checkpoint Recovery
"""

import asyncio


class TestCheckpointRecovery:
    """Test suite for checkpoint-based recovery."""

    def test_recovery_info_dataclass(self):
        """Test RecoveryInfo dataclass."""
        from verlex_server.gateway.backend.recovery import RecoveryInfo, RecoveryState

        info = RecoveryInfo(
            job_id="job123",
            state=RecoveryState.NONE,
        )

        assert info.job_id == "job123"
        assert info.state == RecoveryState.NONE
        assert info.can_retry_recovery
        assert info.recovery_duration_seconds is None

        # Test state transitions
        info.state = RecoveryState.CHECKPOINT_SAVED
        info.checkpoint_id = "ckpt-001"
        assert info.state == RecoveryState.CHECKPOINT_SAVED

        print("✓ RecoveryInfo dataclass works")

    def test_recovery_manager_creation(self):
        """Test CheckpointRecoveryManager creation."""
        from verlex_server.gateway.backend.recovery import CheckpointRecoveryManager

        manager = CheckpointRecoveryManager(
            checkpoint_interval=60,
            max_recovery_attempts=5,
        )

        assert manager.checkpoint_interval == 60
        assert manager.max_recovery_attempts == 5

        print("✓ CheckpointRecoveryManager creation works")

    def test_checkpointer_creation(self):
        """Test getting checkpointer for a job."""
        from verlex_server.gateway.backend.recovery import CheckpointRecoveryManager
        from verlex_server.gateway.core.checkpointing import Framework

        manager = CheckpointRecoveryManager()

        # Get checkpointer (creates if not exists)
        checkpointer = manager.get_checkpointer("job123", Framework.GENERIC)
        assert checkpointer is not None
        assert checkpointer.job_id == "job123"

        # Get same checkpointer again
        checkpointer2 = manager.get_checkpointer("job123")
        assert checkpointer2 is checkpointer

        print("✓ Checkpointer creation works")

    def test_save_and_restore_checkpoint(self):
        """Test saving and restoring checkpoints."""
        from verlex_server.gateway.backend.recovery import CheckpointRecoveryManager, RecoveryState

        async def run_test():
            manager = CheckpointRecoveryManager()
            job_id = "test_job_001"

            # Save a checkpoint
            state = {
                "epoch": 5,
                "step": 1000,
                "model_weights": [1.0, 2.0, 3.0],
                "optimizer_state": {"lr": 0.001},
            }

            metadata = await manager.save_checkpoint(
                job_id=job_id,
                state=state,
                reason="test",
                provider="local",
                region="local",
            )

            assert metadata is not None
            assert metadata.job_id == job_id
            assert metadata.epoch == 5
            assert metadata.step == 1000

            # Check recovery info
            info = manager.get_recovery_info(job_id)
            assert info is not None
            assert info.state == RecoveryState.CHECKPOINT_SAVED
            assert info.checkpoint_id == metadata.checkpoint_id

            # Get latest checkpoint
            latest = await manager.get_latest_checkpoint(job_id)
            assert latest is not None
            assert latest.checkpoint_id == metadata.checkpoint_id

            # Restore checkpoint
            restored_state, restored_meta = await manager.restore_checkpoint(job_id)

            assert restored_state["epoch"] == 5
            assert restored_state["step"] == 1000
            assert restored_state["model_weights"] == [1.0, 2.0, 3.0]

            # Check recovery info updated
            info = manager.get_recovery_info(job_id)
            assert info.state == RecoveryState.RECOVERED

            # Cleanup
            await manager.cleanup_job(job_id, delete_checkpoints=True)

            print("✓ Save and restore checkpoint works")

        asyncio.run(run_test())

    def test_prepare_for_retry(self):
        """Test preparing job for retry with checkpoint."""
        import uuid

        from verlex_server.gateway.backend.recovery import CheckpointRecoveryManager, RecoveryState

        async def run_test():
            manager = CheckpointRecoveryManager()
            # Use unique job ID to avoid conflicts with other tests
            job_id = f"retry_test_job_{uuid.uuid4().hex[:8]}"

            # Cleanup any existing state first
            await manager.cleanup_job(job_id, delete_checkpoints=True)

            # Save a checkpoint first
            await manager.save_checkpoint(
                job_id=job_id,
                state={"step": 500},
                reason="scheduled",
            )

            # Now prepare for retry should find checkpoint
            info = await manager.prepare_for_retry(
                job_id,
                target_provider="gcp",
                target_region="us-central1",
            )

            assert info is not None
            assert info.state == RecoveryState.CHECKPOINT_SAVED
            assert info.target_provider == "gcp"
            assert info.target_region == "us-central1"

            # Cleanup
            await manager.cleanup_job(job_id, delete_checkpoints=True)

            print("✓ Prepare for retry works")

        asyncio.run(run_test())

    def test_interruption_handler(self):
        """Test handling spot interruption."""
        from verlex_server.gateway.backend.recovery import CheckpointRecoveryManager

        async def run_test():
            manager = CheckpointRecoveryManager()
            job_id = "interrupted_job"

            # Simulate interruption with state save
            state = {
                "epoch": 10,
                "step": 5000,
                "important_data": "save_me",
            }

            metadata = await manager.handle_interruption(
                job_id=job_id,
                state=state,
                provider="gcp",
                region="us-central1",
                instance_id="instance-123",
                warning_seconds=30,
            )

            assert metadata is not None

            info = manager.get_recovery_info(job_id)
            assert info.interruption_time is not None
            assert info.source_provider == "gcp"
            assert info.source_instance == "instance-123"

            # Cleanup
            await manager.cleanup_job(job_id, delete_checkpoints=True)

            print("✓ Interruption handler works")

        asyncio.run(run_test())

    def test_callbacks(self):
        """Test recovery manager callbacks."""
        from verlex_server.gateway.backend.recovery import CheckpointRecoveryManager

        async def run_test():
            manager = CheckpointRecoveryManager()
            job_id = "callback_test"

            # Track callback calls
            saved_checkpoints = []
            started_recoveries = []
            completed_recoveries = []

            manager.on_checkpoint_saved(
                lambda jid, meta: saved_checkpoints.append((jid, meta.checkpoint_id))
            )
            manager.on_recovery_started(
                lambda jid: started_recoveries.append(jid)
            )
            manager.on_recovery_completed(
                lambda jid, state: completed_recoveries.append((jid, state))
            )

            # Save checkpoint
            await manager.save_checkpoint(job_id, {"step": 1}, reason="test")
            assert len(saved_checkpoints) == 1
            assert saved_checkpoints[0][0] == job_id

            # Restore checkpoint
            await manager.restore_checkpoint(job_id)
            assert len(started_recoveries) == 1
            assert len(completed_recoveries) == 1
            assert completed_recoveries[0][0] == job_id

            # Cleanup
            await manager.cleanup_job(job_id, delete_checkpoints=True)

            print("✓ Callbacks work")

        asyncio.run(run_test())

    def test_global_instance(self):
        """Test global recovery manager instance."""
        from verlex_server.gateway.backend.recovery import configure_recovery, get_recovery_manager

        # Get default instance
        manager1 = get_recovery_manager()
        assert manager1 is not None

        # Configure with custom settings
        manager2 = configure_recovery(
            checkpoint_interval=120,
            max_recovery_attempts=5,
        )

        assert manager2.checkpoint_interval == 120
        assert manager2.max_recovery_attempts == 5

        # Get should return configured instance
        manager3 = get_recovery_manager()
        assert manager3 is manager2

        print("✓ Global instance works")


def run_all_tests():
    """Run all tests."""
    tests = TestCheckpointRecovery()

    print("Testing Checkpoint Recovery...")
    print("-" * 40)

    tests.test_recovery_info_dataclass()
    tests.test_recovery_manager_creation()
    tests.test_checkpointer_creation()
    tests.test_save_and_restore_checkpoint()
    tests.test_prepare_for_retry()
    tests.test_interruption_handler()
    tests.test_callbacks()
    tests.test_global_instance()

    print("-" * 40)
    print("All recovery tests passed! ✓")


if __name__ == "__main__":
    run_all_tests()
