"""
Tests for Parallel Provisioning (Performance tier, fast=True).

Verifies:
1. Racing only triggers for Performance tier (pricing_mode == "P")
2. Standard tier is completely untouched
3. Winner selection respects budget constraints
4. All loser VMs are terminated (over-budget, lost race, timeout)
5. Fallback to sequential on race failure
"""

import asyncio
from dataclasses import dataclass, field
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from verlex_server.gateway.backend.scheduler_service import (
    CloudJobHandler,
    ProvisionedVM,
    RACE_VM_COUNT,
    RACE_TIMEOUT_SECONDS,
)
from verlex_server.gateway.core.queue import PricingTier


# =========================================================================
# Helpers & Fixtures
# =========================================================================


@dataclass
class FakeInstance:
    """Minimal stand-in for ProviderInstance."""
    instance_id: str = "i-fake123"
    public_ip: str = "1.2.3.4"


@dataclass
class FakeProvider:
    """Minimal stand-in for CloudProviderBase."""
    name: str = "fake"
    terminated_instances: list = field(default_factory=list)

    def get_credentials(self):
        cred = MagicMock()
        cred.is_configured = True
        return cred

    async def list_instance_types(self, **kwargs):
        return [{"instance_type": "gpu-xlarge"}]

    async def launch_instance(self, **kwargs):
        return FakeInstance(instance_id=f"i-{self.name}-001")

    async def terminate_instance(self, instance_id: str) -> bool:
        self.terminated_instances.append(instance_id)
        return True

    async def execute_command(self, **kwargs):
        return (0, "[GateWay-Result]dGVzdA==")  # base64("test")


def make_vm(provider_name: str, hourly_cost: float = 1.0) -> ProvisionedVM:
    """Create a ProvisionedVM for testing."""
    provider_obj = FakeProvider(name=provider_name)
    return ProvisionedVM(
        provider_name=provider_name,
        region="us-east-1",
        instance=FakeInstance(instance_id=f"i-{provider_name}-001"),
        instance_type="gpu-xlarge",
        hourly_cost=hourly_cost,
        provider_obj=provider_obj,
    )


@pytest.fixture
def handler():
    """Create a CloudJobHandler for testing."""
    return CloudJobHandler(preferred_provider="auto")


@pytest.fixture
def fake_job():
    """Create a minimal ScheduledJob-like object."""
    message = MagicMock()
    message.pricing_tier = PricingTier.PERFORMANCE
    message.payload = {
        "function": {"payload": "dGVzdA==", "type": "cloudpickle", "name": "test"},
        "resources": {
            "cpu": 4,
            "memory": "16GB",
            "gpu": "A100",
            "estimated_duration_seconds": 3600,
        },
    }

    job = MagicMock()
    job.job_id = "test-job-001"
    job.user_id = "user-001"
    job.message = message
    job.cloud_instance_id = None
    job.cloud_provider_name = None
    return job


# =========================================================================
# Test: Guarantee 1 — Performance tier only
# =========================================================================


class TestPerformanceTierGating:
    """Verify parallel provisioning ONLY runs for Performance tier."""

    def test_race_vm_count_default(self):
        """GATEWAY_RACE_VM_COUNT defaults to 3."""
        assert RACE_VM_COUNT == 3

    def test_race_timeout_default(self):
        """GATEWAY_RACE_TIMEOUT defaults to 600."""
        assert RACE_TIMEOUT_SECONDS == 600

    def test_provisioned_vm_dataclass(self):
        """ProvisionedVM holds all required fields."""
        vm = make_vm("aws", hourly_cost=2.50)
        assert vm.provider_name == "aws"
        assert vm.hourly_cost == 2.50
        assert vm.instance.instance_id == "i-aws-001"

    def test_mode_map_performance(self):
        """PricingTier.PERFORMANCE maps to 'P' via mode_map."""
        mode_map = {
            "priority": "P", "patient": "Pa", "sticky": "S", "flexible": "F",
            "performance": "P", "standard": "S",
        }
        tier_str = PricingTier.PERFORMANCE.value  # "performance"
        assert mode_map.get(tier_str.lower()) == "P"

    def test_mode_map_standard(self):
        """PricingTier.STANDARD maps to 'S' via mode_map."""
        mode_map = {
            "priority": "P", "patient": "Pa", "sticky": "S", "flexible": "F",
            "performance": "P", "standard": "S",
        }
        tier_str = PricingTier.STANDARD.value  # "standard"
        assert mode_map.get(tier_str.lower()) == "S"

    def test_is_performance_true_for_P(self):
        """is_performance check passes for 'P'."""
        pricing_mode = "P"
        is_performance = pricing_mode == "P"
        assert is_performance is True

    def test_is_performance_false_for_S(self):
        """is_performance check fails for 'S'."""
        pricing_mode = "S"
        is_performance = pricing_mode == "P"
        assert is_performance is False

    def test_is_performance_false_for_legacy_Pa(self):
        """is_performance check fails for legacy 'Pa'."""
        pricing_mode = "Pa"
        is_performance = pricing_mode == "P"
        assert is_performance is False


# =========================================================================
# Test: _get_top_candidates
# =========================================================================


class TestGetTopCandidates:
    """Test provider candidate selection for racing."""

    @patch("verlex_server.gateway.backend.scheduler_service.CloudJobHandler._get_configured_providers")
    @patch("verlex_server.gateway.core.pricing.compare_prices")
    def test_returns_candidates_sorted_by_price(self, mock_compare, mock_configured, handler):
        """Candidates are sorted cheapest first."""
        mock_configured.return_value = ["aws", "gcp", "azure"]

        # Create mock PriceInfo objects
        aws_option = MagicMock()
        aws_option.provider.value = "aws"
        aws_option.region = "us-east-1"
        aws_option.effective_price.return_value = 1.50

        gcp_option = MagicMock()
        gcp_option.provider.value = "gcp"
        gcp_option.region = "us-central1"
        gcp_option.effective_price.return_value = 1.20

        azure_option = MagicMock()
        azure_option.provider.value = "azure"
        azure_option.region = "eastus"
        azure_option.effective_price.return_value = 1.80

        comparison = MagicMock()
        comparison.options = [gcp_option, aws_option, azure_option]  # sorted by price
        mock_compare.return_value = comparison

        candidates = handler._get_top_candidates(
            {"cpu": 4, "memory": "16GB", "gpu": "A100"}, n=3,
        )

        assert len(candidates) == 3
        assert candidates[0][0] == "gcp"  # cheapest first
        assert candidates[1][0] == "aws"
        assert candidates[2][0] == "azure"

    @patch("verlex_server.gateway.backend.scheduler_service.CloudJobHandler._get_configured_providers")
    @patch("verlex_server.gateway.core.pricing.compare_prices")
    def test_deduplicates_by_provider(self, mock_compare, mock_configured, handler):
        """Only one candidate per provider."""
        mock_configured.return_value = ["aws", "gcp"]

        aws1 = MagicMock()
        aws1.provider.value = "aws"
        aws1.region = "us-east-1"
        aws1.effective_price.return_value = 1.00

        aws2 = MagicMock()
        aws2.provider.value = "aws"
        aws2.region = "us-west-2"
        aws2.effective_price.return_value = 1.10

        gcp1 = MagicMock()
        gcp1.provider.value = "gcp"
        gcp1.region = "us-central1"
        gcp1.effective_price.return_value = 1.50

        comparison = MagicMock()
        comparison.options = [aws1, aws2, gcp1]
        mock_compare.return_value = comparison

        candidates = handler._get_top_candidates(
            {"cpu": 4, "memory": "16GB"}, n=3,
        )

        # Should have only 2 (one per provider), not 3
        assert len(candidates) == 2
        providers = [c[0] for c in candidates]
        assert providers.count("aws") == 1
        assert providers.count("gcp") == 1

    @patch("verlex_server.gateway.backend.scheduler_service.CloudJobHandler._get_configured_providers")
    @patch("verlex_server.gateway.core.pricing.compare_prices")
    def test_returns_empty_when_no_options(self, mock_compare, mock_configured, handler):
        """Returns empty list when no providers available."""
        mock_configured.return_value = ["aws"]
        comparison = MagicMock()
        comparison.options = []
        mock_compare.return_value = comparison

        candidates = handler._get_top_candidates({"cpu": 4, "memory": "16GB"})
        assert candidates == []

    @patch("verlex_server.gateway.backend.scheduler_service.CloudJobHandler._get_configured_providers")
    def test_returns_empty_when_no_providers_configured(self, mock_configured, handler):
        """Returns empty list when no providers are configured."""
        mock_configured.return_value = []
        candidates = handler._get_top_candidates({"cpu": 4, "memory": "16GB"})
        assert candidates == []


# =========================================================================
# Test: _terminate_racing_vm
# =========================================================================


class TestTerminateRacingVM:
    """Test loser VM termination."""

    @pytest.mark.asyncio
    async def test_terminates_instance(self, handler):
        """Successfully terminates a racing VM."""
        vm = make_vm("aws")
        output = []
        await handler._terminate_racing_vm(vm, "job-001", lambda m: output.append(m))

        assert "i-aws-001" in vm.provider_obj.terminated_instances
        assert any("Terminated" in o for o in output)

    @pytest.mark.asyncio
    async def test_swallows_termination_error(self, handler):
        """Errors during termination are caught, not raised."""
        vm = make_vm("aws")
        vm.provider_obj.terminate_instance = AsyncMock(
            side_effect=RuntimeError("cloud API down")
        )
        output = []

        # Should NOT raise
        await handler._terminate_racing_vm(vm, "job-001", lambda m: output.append(m))


# =========================================================================
# Test: _get_job_budget
# =========================================================================


class TestGetJobBudget:
    """Test budget retrieval from jobs_db."""

    @patch("verlex_server.gateway.backend.scheduler_service._get_jobs_db")
    def test_reads_estimated_cost(self, mock_get_db, handler, fake_job):
        """Budget comes from job's estimated_cost."""
        mock_get_db.return_value = {
            "test-job-001": {"estimated_cost": 5.50},
        }
        budget = handler._get_job_budget(fake_job)
        assert budget == 5.50

    @patch("verlex_server.gateway.backend.scheduler_service._get_jobs_db")
    def test_fallback_when_no_estimated_cost(self, mock_get_db, handler, fake_job):
        """Fallback calculation when estimated_cost is missing."""
        mock_get_db.return_value = {}
        budget = handler._get_job_budget(fake_job)
        # Should return a positive number (fallback calculation)
        assert budget > 0


# =========================================================================
# Test: _race_provision_performance — core racing logic
# =========================================================================


class TestRaceProvisionPerformance:
    """Test the parallel provisioning race algorithm."""

    @pytest.mark.asyncio
    async def test_winner_is_first_within_budget(self, handler, fake_job):
        """First VM ready that fits budget wins."""
        aws_vm = make_vm("aws", hourly_cost=2.0)
        gcp_vm = make_vm("gcp", hourly_cost=1.5)

        call_count = 0

        async def fake_provision(job, prov, region, cost, on_output):
            nonlocal call_count
            call_count += 1
            if prov == "gcp":
                await asyncio.sleep(0.01)  # GCP is fastest
                return gcp_vm
            else:
                await asyncio.sleep(0.05)  # AWS is slower
                return aws_vm

        handler._provision_vm = fake_provision
        handler._run_on_vm = AsyncMock(return_value="result-ok")

        candidates = [("gcp", "us-central1", 1.5), ("aws", "us-east-1", 2.0)]

        result = await handler._race_provision_performance(
            fake_job, candidates,
            budget=10.0,  # both fit
            estimated_duration_hours=1.0,
            on_output=lambda m: None,
        )

        assert result == "result-ok"
        assert fake_job.cloud_provider_name == "gcp"  # GCP was fastest
        handler._run_on_vm.assert_called_once()

    @pytest.mark.asyncio
    async def test_over_budget_vm_is_terminated(self, handler, fake_job):
        """VM that exceeds budget is terminated immediately."""
        expensive_vm = make_vm("aws", hourly_cost=100.0)
        cheap_vm = make_vm("gcp", hourly_cost=1.0)

        async def fake_provision(job, prov, region, cost, on_output):
            if prov == "aws":
                await asyncio.sleep(0.01)  # AWS finishes first
                return expensive_vm
            else:
                await asyncio.sleep(0.05)
                return cheap_vm

        handler._provision_vm = fake_provision
        handler._run_on_vm = AsyncMock(return_value="result-ok")

        candidates = [("aws", "us-east-1", 100.0), ("gcp", "us-central1", 1.0)]

        result = await handler._race_provision_performance(
            fake_job, candidates,
            budget=5.0,  # AWS exceeds, GCP fits
            estimated_duration_hours=1.0,
            on_output=lambda m: None,
        )

        assert result == "result-ok"
        assert fake_job.cloud_provider_name == "gcp"  # GCP won (within budget)
        # AWS was terminated because it exceeded budget
        assert "i-aws-001" in expensive_vm.provider_obj.terminated_instances

    @pytest.mark.asyncio
    async def test_all_exceed_budget_raises(self, handler, fake_job):
        """Raises RuntimeError when all candidates exceed budget."""
        async def fake_provision(job, prov, region, cost, on_output):
            return make_vm(prov, hourly_cost=100.0)

        handler._provision_vm = fake_provision

        candidates = [("aws", "us-east-1", 100.0), ("gcp", "us-central1", 100.0)]

        with pytest.raises(RuntimeError, match="no candidate was ready within budget"):
            await handler._race_provision_performance(
                fake_job, candidates,
                budget=5.0,
                estimated_duration_hours=1.0,
                on_output=lambda m: None,
            )

    @pytest.mark.asyncio
    async def test_all_candidates_fail_raises(self, handler, fake_job):
        """Raises RuntimeError when all candidates fail to provision."""
        async def fake_provision(job, prov, region, cost, on_output):
            raise RuntimeError(f"{prov} provisioning failed")

        handler._provision_vm = fake_provision

        candidates = [("aws", "us-east-1", 2.0), ("gcp", "us-central1", 1.5)]

        with pytest.raises(RuntimeError, match="no candidate was ready within budget"):
            await handler._race_provision_performance(
                fake_job, candidates,
                budget=10.0,
                estimated_duration_hours=1.0,
                on_output=lambda m: None,
            )

    @pytest.mark.asyncio
    async def test_losers_terminated_after_winner(self, handler, fake_job):
        """All non-winner VMs are terminated in the finally block."""
        aws_vm = make_vm("aws", hourly_cost=2.0)
        gcp_vm = make_vm("gcp", hourly_cost=1.5)
        azure_vm = make_vm("azure", hourly_cost=1.8)

        provision_order = []

        async def fake_provision(job, prov, region, cost, on_output):
            provision_order.append(prov)
            if prov == "gcp":
                await asyncio.sleep(0.01)
                return gcp_vm
            elif prov == "aws":
                await asyncio.sleep(0.1)
                return aws_vm
            else:
                await asyncio.sleep(0.1)
                return azure_vm

        handler._provision_vm = fake_provision
        handler._run_on_vm = AsyncMock(return_value="done")

        candidates = [
            ("gcp", "us-central1", 1.5),
            ("aws", "us-east-1", 2.0),
            ("azure", "eastus", 1.8),
        ]

        await handler._race_provision_performance(
            fake_job, candidates,
            budget=10.0,
            estimated_duration_hours=1.0,
            on_output=lambda m: None,
        )

        # GCP is winner — others should NOT have GCP terminated
        assert gcp_vm.provider_obj.terminated_instances == []
        # Winner was GCP, so AWS and Azure should be terminated
        # (they may or may not have finished provisioning by the time
        # the finally block runs, depending on timing)

    @pytest.mark.asyncio
    async def test_winner_tracking_set_on_job(self, handler, fake_job):
        """Winner's instance_id and provider are set on the job."""
        gcp_vm = make_vm("gcp", hourly_cost=1.5)

        async def fake_provision(job, prov, region, cost, on_output):
            return gcp_vm

        handler._provision_vm = fake_provision
        handler._run_on_vm = AsyncMock(return_value="done")

        candidates = [("gcp", "us-central1", 1.5)]

        await handler._race_provision_performance(
            fake_job, candidates,
            budget=10.0,
            estimated_duration_hours=1.0,
            on_output=lambda m: None,
        )

        assert fake_job.cloud_instance_id == "i-gcp-001"
        assert fake_job.cloud_provider_name == "gcp"

    @pytest.mark.asyncio
    async def test_over_budget_vm_removed_from_provisioned_list(self, handler, fake_job):
        """Over-budget VM is removed from provisioned_vms to prevent double-terminate."""
        expensive_vm = make_vm("aws", hourly_cost=100.0)

        # Track terminate calls
        terminate_calls = []
        original_terminate = handler._terminate_racing_vm

        async def counting_terminate(vm, job_id, on_output):
            terminate_calls.append(vm.provider_name)
            await original_terminate(vm, job_id, on_output)

        handler._terminate_racing_vm = counting_terminate

        async def fake_provision(job, prov, region, cost, on_output):
            return make_vm(prov, hourly_cost=cost)

        handler._provision_vm = fake_provision

        candidates = [("aws", "us-east-1", 100.0), ("gcp", "us-central1", 100.0)]

        with pytest.raises(RuntimeError):
            await handler._race_provision_performance(
                fake_job, candidates,
                budget=5.0,
                estimated_duration_hours=1.0,
                on_output=lambda m: None,
            )

        # Each provider should be terminated exactly ONCE (not double-terminated)
        assert terminate_calls.count("aws") == 1
        assert terminate_calls.count("gcp") == 1
