"""
Tests for Preemption Simulation Module

Tests the spot/preemptible instance preemption simulation functionality.
"""

import asyncio
from datetime import datetime, timedelta
from unittest.mock import patch

import pytest

from verlex_server.gateway.backend.preemption import (
    JobState,
    PreemptionConfig,
    PreemptionEvent,
    PreemptionHandler,
    PreemptionReason,
    PreemptionSimulator,
    PreemptionStrategy,
    configure_simulator,
    get_simulator,
    register_job,
    unregister_job,
)


class TestPreemptionConfig:
    """Tests for PreemptionConfig."""

    def test_default_values(self):
        """Test default configuration values."""
        config = PreemptionConfig()

        assert config.strategy == PreemptionStrategy.RANDOM
        assert config.check_interval_seconds == 30
        assert config.base_probability == 0.05
        assert config.warning_seconds == 120
        assert config.enabled is True

    def test_from_env(self):
        """Test creating config from environment."""
        with patch.dict('os.environ', {
            'GATEWAY_PREEMPTION_STRATEGY': 'time_based',
            'GATEWAY_PREEMPTION_CHECK_INTERVAL': '10',
            'GATEWAY_PREEMPTION_PROBABILITY': '0.1',
            'GATEWAY_PREEMPTION_ENABLED': 'true',
        }):
            config = PreemptionConfig.from_env()

            assert config.strategy == PreemptionStrategy.TIME_BASED
            assert config.check_interval_seconds == 10
            assert config.base_probability == 0.1

    def test_for_testing_default(self):
        """Test config for testing (default)."""
        config = PreemptionConfig.for_testing()

        assert config.strategy == PreemptionStrategy.TIME_BASED
        assert config.min_runtime_seconds == 30
        assert config.max_runtime_seconds == 300
        assert config.enabled is True

    def test_for_testing_aggressive(self):
        """Test config for aggressive testing."""
        config = PreemptionConfig.for_testing(aggressive=True)

        assert config.strategy == PreemptionStrategy.RANDOM
        assert config.base_probability == 0.3
        assert config.check_interval_seconds == 5
        assert config.warning_seconds == 10


class TestJobState:
    """Tests for JobState."""

    def test_runtime_seconds(self):
        """Test runtime calculation."""
        now = datetime.utcnow()
        state = JobState(
            job_id="test",
            container_id=None,
            started_at=now - timedelta(seconds=60),
        )

        runtime = state.runtime_seconds

        # Should be approximately 60 seconds
        assert 59 <= runtime <= 61


class TestPreemptionEvent:
    """Tests for PreemptionEvent."""

    def test_to_dict(self):
        """Test converting event to dictionary."""
        now = datetime.utcnow()
        event = PreemptionEvent(
            job_id="job-123",
            container_id="container-abc",
            timestamp=now,
            reason=PreemptionReason.CAPACITY,
            warning_given=True,
            warning_timestamp=now - timedelta(seconds=120),
            checkpoint_saved=True,
            runtime_seconds=3600,
            metadata={"custom": "value"},
        )

        result = event.to_dict()

        assert result["job_id"] == "job-123"
        assert result["container_id"] == "container-abc"
        assert result["reason"] == "capacity_reclaimed"
        assert result["warning_given"] is True
        assert result["checkpoint_saved"] is True
        assert result["runtime_seconds"] == 3600


class TestPreemptionSimulator:
    """Tests for PreemptionSimulator."""

    @pytest.fixture
    def simulator(self):
        """Create a simulator with testing config."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.TIME_BASED,
            min_runtime_seconds=1,
            max_runtime_seconds=10,
            warning_seconds=2,
            enabled=True,
        )
        return PreemptionSimulator(config)

    def test_register_job(self, simulator):
        """Test job registration."""
        simulator.register_job("job-123", "container-abc")

        assert "job-123" in simulator._jobs
        job = simulator._jobs["job-123"]
        assert job.job_id == "job-123"
        assert job.container_id == "container-abc"

    def test_unregister_job(self, simulator):
        """Test job unregistration."""
        simulator.register_job("job-123")
        assert "job-123" in simulator._jobs

        simulator.unregister_job("job-123")
        assert "job-123" not in simulator._jobs

    @pytest.mark.asyncio
    async def test_check_preemption_not_registered(self, simulator):
        """Test preemption check for unregistered job."""
        preempted, reason = await simulator.check_preemption("unknown-job")

        assert preempted is False
        assert reason is None

    @pytest.mark.asyncio
    async def test_check_preemption_disabled(self):
        """Test preemption check when disabled."""
        config = PreemptionConfig(enabled=False)
        simulator = PreemptionSimulator(config)

        simulator.register_job("job-123")
        preempted, reason = await simulator.check_preemption("job-123")

        assert preempted is False

    @pytest.mark.asyncio
    async def test_check_preemption_before_min_runtime(self, simulator):
        """Test no preemption before minimum runtime."""
        simulator.register_job("job-123")

        # Immediately check - should not preempt
        preempted, reason = await simulator.check_preemption("job-123")

        assert preempted is False

    @pytest.mark.asyncio
    async def test_check_preemption_time_based(self):
        """Test time-based preemption."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.TIME_BASED,
            min_runtime_seconds=0,  # No minimum
            max_runtime_seconds=1,  # Preempt after 1 second
            warning_seconds=0,      # No warning delay
            enabled=True,
        )
        simulator = PreemptionSimulator(config)

        simulator.register_job("job-123")

        # Wait for max runtime
        await asyncio.sleep(1.1)

        # First check schedules preemption
        preempted, reason = await simulator.check_preemption("job-123")

        # Should be preempted now (warning_seconds is 0)
        preempted, reason = await simulator.check_preemption("job-123")

        assert preempted is True
        assert reason == PreemptionReason.SIMULATED

    @pytest.mark.asyncio
    async def test_check_preemption_never_strategy(self):
        """Test never preemption strategy."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.NEVER,
            min_runtime_seconds=0,
            max_runtime_seconds=1,
            enabled=True,
        )
        simulator = PreemptionSimulator(config)

        simulator.register_job("job-123")
        await asyncio.sleep(1.1)

        preempted, reason = await simulator.check_preemption("job-123")

        assert preempted is False

    @pytest.mark.asyncio
    async def test_check_preemption_always_strategy(self):
        """Test always preemption strategy."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.ALWAYS,
            min_runtime_seconds=0,  # No minimum
            warning_seconds=0,      # No warning delay
            enabled=True,
        )
        simulator = PreemptionSimulator(config)

        simulator.register_job("job-123")

        # First check schedules, second executes
        await simulator.check_preemption("job-123")
        preempted, reason = await simulator.check_preemption("job-123")

        assert preempted is True

    @pytest.mark.asyncio
    async def test_force_preempt(self, simulator):
        """Test forced preemption."""
        simulator.register_job("job-123")

        result = await simulator.force_preempt("job-123", PreemptionReason.MANUAL)

        assert result is True

        events = simulator.get_events("job-123")
        assert len(events) == 1
        assert events[0].reason == PreemptionReason.MANUAL

    @pytest.mark.asyncio
    async def test_force_preempt_unknown_job(self, simulator):
        """Test forced preemption of unknown job."""
        result = await simulator.force_preempt("unknown-job")

        assert result is False

    @pytest.mark.asyncio
    async def test_checkpoint_callback(self):
        """Test checkpoint callback on preemption."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.ALWAYS,
            min_runtime_seconds=0,
            warning_seconds=0,
            allow_checkpoint_time=5,
            enabled=True,
        )
        simulator = PreemptionSimulator(config)

        checkpoint_called = False

        async def save_checkpoint():
            nonlocal checkpoint_called
            checkpoint_called = True
            return True

        simulator.register_job("job-123", checkpoint_callback=save_checkpoint)

        # Trigger preemption
        await simulator.check_preemption("job-123")
        await simulator.check_preemption("job-123")

        assert checkpoint_called is True

        events = simulator.get_events("job-123")
        assert len(events) == 1
        assert events[0].checkpoint_saved is True

    def test_get_events(self, simulator):
        """Test getting preemption events."""
        # Initially empty
        events = simulator.get_events()
        assert len(events) == 0

    def test_get_metrics_empty(self, simulator):
        """Test metrics with no events."""
        metrics = simulator.get_metrics()

        assert metrics["total_preemptions"] == 0
        assert metrics["avg_runtime_seconds"] == 0

    @pytest.mark.asyncio
    async def test_get_metrics_with_events(self):
        """Test metrics with events."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.ALWAYS,
            min_runtime_seconds=0,
            warning_seconds=0,
            enabled=True,
        )
        simulator = PreemptionSimulator(config)

        # Create some preemption events
        for i in range(3):
            simulator.register_job(f"job-{i}")
            await simulator.check_preemption(f"job-{i}")
            await simulator.check_preemption(f"job-{i}")

        metrics = simulator.get_metrics()

        assert metrics["total_preemptions"] == 3

    @pytest.mark.asyncio
    async def test_preemptible_context(self, simulator):
        """Test preemptible context manager."""
        checkpoint_called = False

        async def save_checkpoint():
            nonlocal checkpoint_called
            checkpoint_called = True
            return True

        async with simulator.preemptible_context(
            "job-ctx",
            container_id="container-123",
            on_preempt=save_checkpoint,
        ):
            assert "job-ctx" in simulator._jobs

        # Job should be unregistered after context
        assert "job-ctx" not in simulator._jobs

    @pytest.mark.asyncio
    async def test_event_callbacks(self):
        """Test preemption event callbacks."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.ALWAYS,
            min_runtime_seconds=0,
            warning_seconds=0,
            enabled=True,
        )
        simulator = PreemptionSimulator(config)

        callback_events = []

        async def on_event(event):
            callback_events.append(event)

        simulator.add_event_callback(on_event)
        simulator.register_job("job-123")

        await simulator.check_preemption("job-123")
        await simulator.check_preemption("job-123")

        assert len(callback_events) == 1
        assert callback_events[0].job_id == "job-123"


class TestPreemptionHandler:
    """Tests for PreemptionHandler."""

    @pytest.mark.asyncio
    async def test_reschedule_on_preemption(self):
        """Test job rescheduling on preemption."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.ALWAYS,
            min_runtime_seconds=0,
            warning_seconds=0,
            enabled=True,
        )
        simulator = PreemptionSimulator(config)
        handler = PreemptionHandler(simulator)

        rescheduled_jobs = []

        async def reschedule(job_id):
            rescheduled_jobs.append(job_id)

        handler.set_reschedule_callback(reschedule)

        # Set up a job with checkpoint
        async def save_checkpoint():
            return True

        simulator.register_job("job-123", checkpoint_callback=save_checkpoint)

        # Trigger preemption
        await simulator.check_preemption("job-123")
        await simulator.check_preemption("job-123")

        # Give time for callback
        await asyncio.sleep(0.1)

        assert "job-123" in rescheduled_jobs


class TestConvenienceFunctions:
    """Tests for module-level convenience functions."""

    def test_get_simulator(self):
        """Test getting global simulator."""
        simulator = get_simulator()

        assert simulator is not None
        assert isinstance(simulator, PreemptionSimulator)

    def test_configure_simulator(self):
        """Test configuring global simulator."""
        config = PreemptionConfig(strategy=PreemptionStrategy.NEVER)

        simulator = configure_simulator(config)

        assert simulator.config.strategy == PreemptionStrategy.NEVER

    def test_register_unregister_job(self):
        """Test register/unregister via convenience functions."""
        register_job("test-job")

        simulator = get_simulator()
        assert "test-job" in simulator._jobs

        unregister_job("test-job")
        assert "test-job" not in simulator._jobs


class TestPreemptionMonitor:
    """Tests for background preemption monitor."""

    @pytest.mark.asyncio
    async def test_run_preemption_monitor_stopped(self):
        """Test preemption monitor stops on event."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.NEVER,
            check_interval_seconds=1,
            enabled=True,
        )
        simulator = PreemptionSimulator(config)

        simulator.register_job("job-123")
        stop_event = asyncio.Event()

        # Start monitor in background
        monitor_task = asyncio.create_task(
            simulator.run_preemption_monitor("job-123", stop_event)
        )

        # Stop after a brief delay
        await asyncio.sleep(0.1)
        stop_event.set()

        reason = await monitor_task

        assert reason is None  # Stopped normally, not preempted

    @pytest.mark.asyncio
    async def test_run_preemption_monitor_preempted(self):
        """Test preemption monitor returns reason on preemption."""
        config = PreemptionConfig(
            strategy=PreemptionStrategy.TIME_BASED,
            min_runtime_seconds=0,
            max_runtime_seconds=0,  # Immediate preemption
            warning_seconds=0,
            check_interval_seconds=1,
            enabled=True,
        )
        simulator = PreemptionSimulator(config)

        simulator.register_job("job-456")

        # Run monitor
        reason = await asyncio.wait_for(
            simulator.run_preemption_monitor("job-456"),
            timeout=5,
        )

        assert reason == PreemptionReason.SIMULATED
