"""Tests for multi-stage pipeline scheduling (Issue #36)."""

from __future__ import annotations

import pytest

from sagellm_backend.pipeline import MultiStagePipeline, create_multi_stage_pipeline
from sagellm_backend.providers.cpu import CPUBackendProvider


class TestMultiStagePipeline:
    """Validate generic staged pipeline behavior."""

    def test_stage_count_bounds(self) -> None:
        """Pipeline must enforce stage count in [2, 4]."""
        provider = CPUBackendProvider()

        with pytest.raises(ValueError, match=r"stage_count must be in \[2, 4\]"):
            MultiStagePipeline(provider, stage_count=1)

        with pytest.raises(ValueError, match=r"stage_count must be in \[2, 4\]"):
            MultiStagePipeline(provider, stage_count=5)

    def test_factory_and_stage_count(self) -> None:
        """Factory should create pipeline with expected stage count."""
        provider = CPUBackendProvider()
        pipeline = create_multi_stage_pipeline(provider, stage_count=3)
        assert pipeline.stage_count == 3

    def test_double_buffering_stage_rotation(self) -> None:
        """2-stage pipeline should rotate stage index 0/1 in ring order."""
        provider = CPUBackendProvider()
        pipeline = MultiStagePipeline(provider, stage_count=2)
        seen_stage_indices: list[int] = []

        def prefetch_fn(item: int, stage_index: int) -> int:
            seen_stage_indices.append(stage_index)
            return item + stage_index

        def compute_fn(prefetched: int, stage_index: int) -> int:
            return prefetched * (stage_index + 1)

        output = pipeline.run([1, 2, 3, 4], prefetch_fn, compute_fn)

        assert seen_stage_indices == [0, 1, 0, 1]
        assert len(output) == 4

    def test_three_stage_pipeline_support(self) -> None:
        """Pipeline must support 3-stage scheduling as required by issue."""
        provider = CPUBackendProvider()
        pipeline = MultiStagePipeline(provider, stage_count=3)
        stage_used: list[int] = []

        def prefetch_fn(item: int, stage_index: int) -> int:
            stage_used.append(stage_index)
            return item

        def compute_fn(prefetched: int, stage_index: int) -> int:
            return prefetched + stage_index

        output = pipeline.run([10, 10, 10, 10, 10], prefetch_fn, compute_fn)

        assert stage_used == [0, 1, 2, 0, 1]
        assert output == [10, 11, 12, 10, 11]

    def test_metrics_available_after_run(self) -> None:
        """Pipeline should expose run metrics for observability."""
        provider = CPUBackendProvider()
        pipeline = MultiStagePipeline(provider, stage_count=2)

        def prefetch_fn(item: int, stage_index: int) -> int:
            return item + stage_index

        def compute_fn(prefetched: int, stage_index: int) -> int:
            return prefetched

        pipeline.run([1, 2, 3], prefetch_fn, compute_fn)
        metrics = pipeline.last_metrics

        assert metrics is not None
        assert metrics.stage_count == 2
        assert metrics.total_items == 3
        assert set(metrics.stage_latencies_ms.keys()) == {0, 1}
        assert metrics.total_wall_time_ms >= 0.0
        assert 0.0 <= metrics.overlap_ratio <= 1.0
