"""
Tests for elastic resource scaling.

Tests the ability to dynamically scale CPU and memory during job execution
without restarts or interruptions.
"""

import asyncio
import os
from datetime import datetime
from unittest.mock import Mock

import pytest

from verlex_server.gateway.backend.elastic_resources import (
    ElasticConfig,
    ElasticExecutionWrapper,
    ProcessMonitor,
    ResourceLimits,
    ResourceMonitor,
    ResourceUsage,
    ScalingDirection,
    ScalingEvent,
    ScalingReason,
)


class TestResourceLimits:
    """Test ResourceLimits class."""

    def test_can_scale_up_cpu(self):
        """Should allow scaling up when below max."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_cpu_cores=8)
        assert limits.can_scale_up_cpu() is True

    def test_cannot_scale_up_cpu_at_max(self):
        """Should not allow scaling when at max."""
        limits = ResourceLimits(cpu_cores=8, memory_mb=4096, max_cpu_cores=8)
        assert limits.can_scale_up_cpu() is False

    def test_can_scale_down_cpu(self):
        """Should allow scaling down when above min."""
        limits = ResourceLimits(cpu_cores=4, memory_mb=4096, min_cpu_cores=1)
        assert limits.can_scale_down_cpu() is True

    def test_cannot_scale_down_cpu_at_min(self):
        """Should not allow scaling down when at min."""
        limits = ResourceLimits(cpu_cores=1, memory_mb=4096, min_cpu_cores=1)
        assert limits.can_scale_down_cpu() is False

    def test_next_cpu_tier_up(self):
        """Should return next higher CPU tier."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_cpu_cores=32)
        assert limits.next_cpu_tier(ScalingDirection.UP) == 4

    def test_next_cpu_tier_down(self):
        """Should return next lower CPU tier."""
        limits = ResourceLimits(cpu_cores=4, memory_mb=4096, min_cpu_cores=1)
        assert limits.next_cpu_tier(ScalingDirection.DOWN) == 2

    def test_next_cpu_tier_respects_max(self):
        """Should not exceed max when scaling up."""
        limits = ResourceLimits(cpu_cores=4, memory_mb=4096, max_cpu_cores=6)
        # Next tier would be 8, but max is 6
        assert limits.next_cpu_tier(ScalingDirection.UP) == 6

    def test_next_memory_tier_up(self):
        """Should return next higher memory tier."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_memory_mb=65536)
        assert limits.next_memory_tier(ScalingDirection.UP) == 8192

    def test_next_memory_tier_down(self):
        """Should return next lower memory tier."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=8192, min_memory_mb=512)
        assert limits.next_memory_tier(ScalingDirection.DOWN) == 4096


class TestElasticConfig:
    """Test ElasticConfig defaults."""

    def test_default_thresholds(self):
        """Default thresholds should be sensible."""
        config = ElasticConfig()

        # Scale up at 85% CPU usage
        assert config.cpu_scale_up_threshold == 0.85

        # Scale down at 30% CPU usage
        assert config.cpu_scale_down_threshold == 0.30

        # Require sustained usage
        assert config.sustained_checks_required == 3

    def test_aggressive_mode(self):
        """Aggressive mode should be disabled by default."""
        config = ElasticConfig()
        assert config.aggressive_mode is False


class TestScalingEvent:
    """Test ScalingEvent recording."""

    def test_to_dict(self):
        """Should convert to dictionary correctly."""
        event = ScalingEvent(
            timestamp=datetime(2026, 1, 9, 12, 0, 0),
            job_id="test-job",
            resource_type="cpu",
            direction=ScalingDirection.UP,
            reason=ScalingReason.CPU_SATURATION,
            old_value=2.0,
            new_value=4.0,
            usage_at_time=92.5,
        )

        d = event.to_dict()

        assert d["job_id"] == "test-job"
        assert d["resource_type"] == "cpu"
        assert d["direction"] == "up"
        assert d["reason"] == "cpu_saturation"
        assert d["old_value"] == 2.0
        assert d["new_value"] == 4.0
        assert d["usage_at_time"] == 92.5


class TestResourceMonitor:
    """Test ResourceMonitor scaling logic."""

    @pytest.fixture
    def mock_proc_monitor(self):
        """Create a mock ProcessMonitor."""
        monitor = Mock(spec=ProcessMonitor)
        monitor.get_cpu_percent.return_value = 50.0
        monitor.get_memory_usage.return_value = (1024 * 1024 * 1024, 2 * 1024 * 1024 * 1024)
        monitor.get_thread_count.return_value = 4
        return monitor

    def test_should_scale_up_cpu_when_sustained(self):
        """Should scale up when CPU is sustained above threshold."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_cpu_cores=8)
        config = ElasticConfig(sustained_checks_required=3)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
            config=config,
        )

        # Simulate sustained high usage (85% of 2 cores = 170%)
        config.cpu_scale_up_threshold * 100 * limits.cpu_cores  # 170

        monitor._cpu_history = [180, 185, 190]  # All above threshold

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=190,
            cpu_cores_used=1.9,
            memory_used_mb=2048,
            memory_percent=50,
        )

        assert monitor._should_scale_up_cpu(usage) is True

    def test_should_not_scale_up_cpu_when_at_max(self):
        """Should not try to scale when already at max."""
        limits = ResourceLimits(cpu_cores=8, memory_mb=4096, max_cpu_cores=8)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
        )

        monitor._cpu_history = [800, 800, 800]  # 100% of 8 cores

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=800,
            cpu_cores_used=8.0,
            memory_used_mb=2048,
            memory_percent=50,
        )

        # Can't scale up, already at max
        assert monitor._should_scale_up_cpu(usage) is False

    def test_should_not_scale_without_sustained_usage(self):
        """Should not scale on a single spike."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_cpu_cores=8)
        config = ElasticConfig(sustained_checks_required=3)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
            config=config,
        )

        # Only 1 high reading, not sustained
        monitor._cpu_history = [190]

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=190,
            cpu_cores_used=1.9,
            memory_used_mb=2048,
            memory_percent=50,
        )

        assert monitor._should_scale_up_cpu(usage) is False

    def test_should_scale_up_memory_when_sustained(self):
        """Should scale up memory when sustained above threshold."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_memory_mb=32768)
        config = ElasticConfig(
            memory_scale_up_threshold=0.80,
            sustained_checks_required=3,
        )

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
            config=config,
        )

        # All above 80% threshold
        monitor._memory_history = [85, 88, 90]

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=50,
            cpu_cores_used=0.5,
            memory_used_mb=3600,
            memory_percent=88,
        )

        assert monitor._should_scale_up_memory(usage) is True

    def test_should_scale_down_cpu_when_sustained_low(self):
        """Should scale down when CPU is sustained below threshold."""
        limits = ResourceLimits(cpu_cores=4, memory_mb=4096, min_cpu_cores=1)
        config = ElasticConfig(
            cpu_scale_down_threshold=0.30,
            sustained_checks_required=3,
            enable_scale_down=True,
        )

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
            config=config,
        )

        # Need 3x sustained checks for scale down
        # 30% of 4 cores = 120% threshold
        # Usage below that for 9 consecutive readings
        monitor._cpu_history = [50, 55, 45, 60, 50, 55, 48, 52, 50]

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=50,
            cpu_cores_used=0.5,
            memory_used_mb=2048,
            memory_percent=50,
        )

        assert monitor._should_scale_down_cpu(usage) is True

    def test_scale_down_disabled(self):
        """Should not scale down when disabled."""
        limits = ResourceLimits(cpu_cores=4, memory_mb=4096, min_cpu_cores=1)
        config = ElasticConfig(enable_scale_down=False)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
            config=config,
        )

        monitor._cpu_history = [10, 10, 10, 10, 10, 10, 10, 10, 10]

        ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=10,
            cpu_cores_used=0.1,
            memory_used_mb=1024,
            memory_percent=25,
        )

        # Config has scale down disabled
        assert monitor.config.enable_scale_down is False


class TestResourceMonitorScaling:
    """Test actual scaling operations."""

    @pytest.mark.asyncio
    async def test_scale_cpu_up(self):
        """Should scale CPU up and record event."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_cpu_cores=8)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
        )

        # Mock cgroup as unavailable (no actual cgroup changes)
        monitor.cgroup_available = False

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=190,
            cpu_cores_used=1.9,
            memory_used_mb=2048,
            memory_percent=50,
        )

        await monitor._scale_cpu(ScalingDirection.UP, usage, ScalingReason.CPU_SATURATION)

        # Check limits were updated
        assert monitor.limits.cpu_cores == 4  # 2 -> 4

        # Check event was recorded
        assert len(monitor.scaling_events) == 1
        event = monitor.scaling_events[0]
        assert event.direction == ScalingDirection.UP
        assert event.old_value == 2
        assert event.new_value == 4
        assert event.reason == ScalingReason.CPU_SATURATION

    @pytest.mark.asyncio
    async def test_scale_memory_up(self):
        """Should scale memory up and record event."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_memory_mb=32768)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
        )

        monitor.cgroup_available = False

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=100,
            cpu_cores_used=1.0,
            memory_used_mb=3500,
            memory_percent=85,
        )

        await monitor._scale_memory(ScalingDirection.UP, usage, ScalingReason.MEMORY_PRESSURE)

        assert monitor.limits.memory_mb == 8192  # 4096 -> 8192
        assert len(monitor.scaling_events) == 1
        assert monitor.scaling_events[0].resource_type == "memory"

    @pytest.mark.asyncio
    async def test_on_scale_callback(self):
        """Should call callback when scaling."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_cpu_cores=8)

        callback_events = []

        def on_scale(event):
            callback_events.append(event)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
            on_scale=on_scale,
        )

        monitor.cgroup_available = False

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=190,
            cpu_cores_used=1.9,
            memory_used_mb=2048,
            memory_percent=50,
        )

        await monitor._scale_cpu(ScalingDirection.UP, usage, ScalingReason.CPU_SATURATION)

        assert len(callback_events) == 1
        assert callback_events[0].new_value == 4


class TestBillingCalculation:
    """Test usage-based billing calculations."""

    def test_billing_usage_calculation(self):
        """Should calculate CPU-seconds and memory-GB-seconds."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096)
        config = ElasticConfig(check_interval_seconds=1.0)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
            config=config,
        )

        # Simulate 10 samples of usage
        for _i in range(10):
            monitor.usage_samples.append(ResourceUsage(
                timestamp=datetime.utcnow(),
                cpu_percent=100,  # 1 core
                cpu_cores_used=1.0,
                memory_used_mb=2048,  # 2 GB
                memory_percent=50,
            ))

        billing = monitor.get_billing_usage()

        # 10 samples * 1.0 core * 1.0 second = 10 CPU-seconds
        assert billing["cpu_seconds"] == 10.0

        # 10 samples * 2 GB * 1.0 second = 20 GB-seconds
        assert billing["memory_gb_seconds"] == 20.0

    def test_billing_with_varying_usage(self):
        """Should handle varying usage correctly."""
        limits = ResourceLimits(cpu_cores=4, memory_mb=8192)
        config = ElasticConfig(check_interval_seconds=1.0)

        monitor = ResourceMonitor(
            job_id="test",
            pid=1234,
            initial_limits=limits,
            config=config,
        )

        # 5 samples at low usage
        for _i in range(5):
            monitor.usage_samples.append(ResourceUsage(
                timestamp=datetime.utcnow(),
                cpu_percent=100,  # 1 core
                cpu_cores_used=1.0,
                memory_used_mb=1024,  # 1 GB
                memory_percent=12.5,
            ))

        # 5 samples at high usage
        for _i in range(5):
            monitor.usage_samples.append(ResourceUsage(
                timestamp=datetime.utcnow(),
                cpu_percent=300,  # 3 cores
                cpu_cores_used=3.0,
                memory_used_mb=4096,  # 4 GB
                memory_percent=50,
            ))

        billing = monitor.get_billing_usage()

        # 5 * 1.0 + 5 * 3.0 = 20 CPU-seconds
        assert billing["cpu_seconds"] == 20.0

        # 5 * 1.0 + 5 * 4.0 = 25 GB-seconds
        assert billing["memory_gb_seconds"] == 25.0


class TestElasticExecutionWrapper:
    """Test the execution wrapper."""

    def test_wrapper_initialization(self):
        """Should initialize with correct limits."""
        wrapper = ElasticExecutionWrapper(
            job_id="test-job",
            initial_cpu=2,
            initial_memory_mb=4096,
            max_cpu=16,
            max_memory_mb=65536,
        )

        assert wrapper.limits.cpu_cores == 2
        assert wrapper.limits.memory_mb == 4096
        assert wrapper.limits.max_cpu_cores == 16
        assert wrapper.limits.max_memory_mb == 65536

    @pytest.mark.asyncio
    async def test_run_async_monitors_function(self):
        """Should monitor an async function."""
        wrapper = ElasticExecutionWrapper(
            job_id="test-job",
            initial_cpu=2,
            initial_memory_mb=4096,
        )

        async def simple_task():
            await asyncio.sleep(0.1)
            return 42

        result = await wrapper.run_async(simple_task)

        assert result == 42
        assert wrapper.monitor is not None


class TestCgroupController:
    """Test cgroup controller (mocked for CI environments)."""

    def test_cpu_limit_format(self):
        """Test CPU limit calculation."""
        # For 2 cores with 100ms period: quota = 200000 (200ms)
        period = 100000
        cores = 2.0
        quota = int(cores * period)

        assert quota == 200000
        assert f"{quota} {period}" == "200000 100000"

    def test_memory_limit_format(self):
        """Test memory limit calculation."""
        memory_mb = 4096
        memory_bytes = memory_mb * 1024 * 1024

        assert memory_bytes == 4294967296  # 4GB in bytes


class TestIntegration:
    """Integration tests for elastic resources."""

    @pytest.mark.asyncio
    async def test_full_scaling_scenario(self):
        """Test a complete scaling scenario."""
        limits = ResourceLimits(
            cpu_cores=2,
            memory_mb=4096,
            max_cpu_cores=8,
            max_memory_mb=32768,
        )
        config = ElasticConfig(
            sustained_checks_required=2,
            check_interval_seconds=0.1,
        )

        scaling_events = []

        monitor = ResourceMonitor(
            job_id="integration-test",
            pid=os.getpid(),
            initial_limits=limits,
            config=config,
            on_scale=lambda e: scaling_events.append(e),
        )

        # Don't actually use cgroups
        monitor.cgroup_available = False

        # Manually add high usage history
        monitor._cpu_history = [180, 185]  # Above 85% of 2 cores (170)

        usage = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=190,
            cpu_cores_used=1.9,
            memory_used_mb=2048,
            memory_percent=50,
        )

        # Check and scale
        await monitor._check_and_scale(usage)

        # Should have scaled up
        assert monitor.limits.cpu_cores == 4
        assert len(scaling_events) == 1

        # Simulate another round with still high usage
        monitor._cpu_history = [380, 390]  # Above 85% of 4 cores (340)

        usage2 = ResourceUsage(
            timestamp=datetime.utcnow(),
            cpu_percent=395,
            cpu_cores_used=3.95,
            memory_used_mb=2048,
            memory_percent=50,
        )

        # Bypass cooldown for test
        monitor._last_scale_up = 0

        await monitor._check_and_scale(usage2)

        # Should have scaled up again
        assert monitor.limits.cpu_cores == 8
        assert len(scaling_events) == 2

    @pytest.mark.asyncio
    async def test_summary_output(self):
        """Test summary includes all relevant info."""
        limits = ResourceLimits(cpu_cores=2, memory_mb=4096, max_cpu_cores=8)

        monitor = ResourceMonitor(
            job_id="summary-test",
            pid=12345,
            initial_limits=limits,
        )

        monitor.cgroup_available = False

        # Add a scaling event
        monitor.scaling_events.append(ScalingEvent(
            timestamp=datetime.utcnow(),
            job_id="summary-test",
            resource_type="cpu",
            direction=ScalingDirection.UP,
            reason=ScalingReason.CPU_SATURATION,
            old_value=2,
            new_value=4,
            usage_at_time=92,
        ))
        monitor.limits.cpu_cores = 4

        summary = monitor.get_summary()

        assert summary["job_id"] == "summary-test"
        assert summary["current_limits"]["cpu_cores"] == 4
        assert summary["max_limits"]["cpu_cores"] == 8
        assert len(summary["scaling_events"]) == 1
        assert "billing" in summary
