"""Tests for Task1.8: Multi-stream Overlap – StreamPlacement + MultiStreamPipeline.

Coverage:
- StreamPlacement construction, properties, helpers
- PipelineStage construction + validation
- MultiStreamPipeline: add_stage, duplicate names, bad dependencies
- Topological ordering (linear chain, fan-in, independent stages)
- Cycle detection
- Pipeline execution: N-stage, dependency ordering, failure propagation
- Placement metadata in PipelineResult
- Integration: 3-stream compute-comm-decode scenario
"""

from __future__ import annotations

import time

import pytest

from sagellm_comm.overlap import (
    MultiStreamPipeline,
    PipelineStage,
    StageResult,
    StreamPlacement,
    create_cpu_stream,
)
from sagellm_comm.overlap.placement import DeviceType

# =============================================================================
# StreamPlacement
# =============================================================================


class TestStreamPlacement:
    """Tests for StreamPlacement."""

    def test_default_placement_is_local_cpu(self) -> None:
        p = StreamPlacement()
        assert p.device_type == DeviceType.CPU
        assert p.device_id == 0
        assert p.rank is None
        assert p.node_id is None
        assert p.is_local

    def test_cpu_factory(self) -> None:
        p = StreamPlacement.cpu(rank=0, node_id="node0")
        assert p.device_type == DeviceType.CPU
        assert p.rank == 0
        assert p.node_id == "node0"
        assert not p.is_local

    def test_cuda_factory(self) -> None:
        p = StreamPlacement.cuda(device_id=1, rank=3)
        assert p.device_type == DeviceType.CUDA
        assert p.device_id == 1
        assert p.rank == 3

    def test_npu_factory(self) -> None:
        p = StreamPlacement.npu(device_id=2)
        assert p.device_type == DeviceType.NPU
        assert p.device_id == 2

    def test_key_format_local(self) -> None:
        p = StreamPlacement.cpu()
        assert p.key == "local/cpu:0"

    def test_key_format_remote(self) -> None:
        p = StreamPlacement.cuda(device_id=0, node_id="node42")
        assert p.key == "node42/cuda:0"

    def test_same_device(self) -> None:
        p0 = StreamPlacement.cuda(device_id=0)
        p0b = StreamPlacement.cuda(device_id=0)
        p1 = StreamPlacement.cuda(device_id=1)
        assert p0.same_device(p0b)
        assert not p0.same_device(p1)

    def test_same_node_local(self) -> None:
        p0 = StreamPlacement.cpu()
        p1 = StreamPlacement.cuda()
        assert p0.same_node(p1)  # Both node_id=None

    def test_same_node_remote(self) -> None:
        p0 = StreamPlacement(node_id="node0")
        p1 = StreamPlacement(node_id="node0")
        p2 = StreamPlacement(node_id="node1")
        assert p0.same_node(p1)
        assert not p0.same_node(p2)

    def test_frozen_immutable(self) -> None:
        p = StreamPlacement.cpu()
        with pytest.raises((AttributeError, TypeError)):
            p.device_id = 99  # type: ignore[misc]

    def test_repr_contains_device(self) -> None:
        p = StreamPlacement.cuda(device_id=2, rank=1, node_id="n0")
        r = repr(p)
        assert "cuda:2" in r
        assert "rank=1" in r
        assert "node=n0" in r


# =============================================================================
# PipelineStage
# =============================================================================


class TestPipelineStage:
    """Tests for PipelineStage construction."""

    def test_construction(self) -> None:
        stream = create_cpu_stream()
        placement = StreamPlacement.cpu()
        stage = PipelineStage(
            name="s1",
            stream=stream,
            placement=placement,
            fn=lambda: None,
        )
        assert stage.name == "s1"
        assert stage.dependencies == []

    def test_empty_name_raises(self) -> None:
        with pytest.raises(ValueError, match="name must not be empty"):
            PipelineStage(
                name="",
                stream=create_cpu_stream(),
                placement=StreamPlacement.cpu(),
                fn=lambda: None,
            )


# =============================================================================
# MultiStreamPipeline – registration & topology
# =============================================================================


class TestMultiStreamPipelineRegistration:
    """Tests for stage registration and topology helpers."""

    def _pipeline(self) -> MultiStreamPipeline:
        return MultiStreamPipeline(trace_id="test-trace")

    def test_add_stage_returns_stage(self) -> None:
        p = self._pipeline()
        stage = p.add_stage("a", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        assert isinstance(stage, PipelineStage)
        assert stage.name == "a"

    def test_stage_count(self) -> None:
        p = self._pipeline()
        assert p.stage_count == 0
        p.add_stage("x", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        assert p.stage_count == 1
        p.add_stage("y", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        assert p.stage_count == 2

    def test_duplicate_name_raises(self) -> None:
        p = self._pipeline()
        p.add_stage("dup", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        with pytest.raises(ValueError, match="already registered"):
            p.add_stage("dup", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)

    def test_missing_dependency_raises_on_topology(self) -> None:
        p = self._pipeline()
        p.add_stage("a", create_cpu_stream(), StreamPlacement.cpu(), lambda: None, ["ghost"])
        with pytest.raises(ValueError, match="unknown stage"):
            p.get_topology_order()

    def test_topology_linear(self) -> None:
        p = self._pipeline()
        p.add_stage("c", create_cpu_stream(), StreamPlacement.cpu(), lambda: None, ["b"])
        p.add_stage("b", create_cpu_stream(), StreamPlacement.cpu(), lambda: None, ["a"])
        p.add_stage("a", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        order = p.get_topology_order()
        assert order.index("a") < order.index("b") < order.index("c")

    def test_topology_fan_in(self) -> None:
        """'c' depends on both 'a' and 'b'."""
        p = self._pipeline()
        p.add_stage("a", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        p.add_stage("b", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        p.add_stage("c", create_cpu_stream(), StreamPlacement.cpu(), lambda: None, ["a", "b"])
        order = p.get_topology_order()
        assert order.index("a") < order.index("c")
        assert order.index("b") < order.index("c")

    def test_topology_independent_stages(self) -> None:
        """Independent stages can appear in any order."""
        p = self._pipeline()
        for name in ("x", "y", "z"):
            p.add_stage(name, create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        order = p.get_topology_order()
        assert set(order) == {"x", "y", "z"}

    def test_cycle_detected(self) -> None:
        p = self._pipeline()
        p.add_stage("a", create_cpu_stream(), StreamPlacement.cpu(), lambda: None, ["c"])
        p.add_stage("b", create_cpu_stream(), StreamPlacement.cpu(), lambda: None, ["a"])
        p.add_stage("c", create_cpu_stream(), StreamPlacement.cpu(), lambda: None, ["b"])
        with pytest.raises(RuntimeError, match="Cycle detected"):
            p.get_topology_order()

    def test_get_placements(self) -> None:
        p = self._pipeline()
        pa = StreamPlacement.cpu(rank=0)
        pb = StreamPlacement.cuda(device_id=0)
        p.add_stage("a", create_cpu_stream(), pa, lambda: None)
        p.add_stage("b", create_cpu_stream(), pb, lambda: None)
        placements = p.get_placements()
        assert placements["a"] == pa
        assert placements["b"] == pb

    def test_empty_pipeline_execute_raises(self) -> None:
        p = self._pipeline()
        with pytest.raises(ValueError, match="empty pipeline"):
            p.execute()


# =============================================================================
# MultiStreamPipeline – execution
# =============================================================================


class TestMultiStreamPipelineExecution:
    """Tests for pipeline execution correctness and ordering."""

    def test_single_stage(self) -> None:
        p = MultiStreamPipeline(trace_id="single")
        results: list[str] = []
        p.add_stage("a", create_cpu_stream(), StreamPlacement.cpu(), lambda: results.append("a"))
        r = p.execute()
        assert r.success
        assert len(r.stage_results) == 1
        assert r.stage_results[0].stage_name == "a"
        assert r.stage_results[0].success
        assert results == ["a"]

    def test_three_stage_linear_ordering(self) -> None:
        """Stages must execute in dependency order even if added out-of-order."""
        p = MultiStreamPipeline(trace_id="linear-3")
        execution_order: list[str] = []
        lock = __import__("threading").Lock()

        def make_fn(name: str, delay: float = 0.0):
            def fn() -> None:
                time.sleep(delay)
                with lock:
                    execution_order.append(name)

            return fn

        # Add in reverse order
        p.add_stage("c", create_cpu_stream(), StreamPlacement.cpu(), make_fn("c"), ["b"])
        p.add_stage("b", create_cpu_stream(), StreamPlacement.cpu(), make_fn("b", 0.02), ["a"])
        p.add_stage("a", create_cpu_stream(), StreamPlacement.cpu(), make_fn("a"))

        r = p.execute()

        assert r.success
        # a must come before b; b must come before c
        assert execution_order.index("a") < execution_order.index("b")
        assert execution_order.index("b") < execution_order.index("c")

    def test_independent_stages_run_concurrently(self) -> None:
        """Three independent stages should run in parallel (total < sum of delays)."""
        p = MultiStreamPipeline(trace_id="parallel-3")
        delay = 0.05

        p.add_stage("x", create_cpu_stream(), StreamPlacement.cpu(), lambda: time.sleep(delay))
        p.add_stage("y", create_cpu_stream(), StreamPlacement.cpu(), lambda: time.sleep(delay))
        p.add_stage("z", create_cpu_stream(), StreamPlacement.cpu(), lambda: time.sleep(delay))

        start = time.perf_counter()
        r = p.execute()
        elapsed = time.perf_counter() - start

        assert r.success
        # Three stages with 50ms each, if sequential = ~150ms; parallel ~ 50ms
        assert elapsed < delay * 2.5, f"Stages did not run concurrently (elapsed={elapsed:.3f}s)"

    def test_five_stage_pipeline(self) -> None:
        """Ensure >2-stream pipeline works correctly (core of Task1.8)."""
        p = MultiStreamPipeline(trace_id="5-stage")
        log: list[str] = []
        lock = __import__("threading").Lock()

        def make_fn(name: str):
            def fn() -> None:
                with lock:
                    log.append(name)

            return fn

        # Linear chain: a -> b -> c (with d and e independent branches)
        # a -> b -> c
        # a -> d
        p.add_stage("a", create_cpu_stream(), StreamPlacement.cpu(rank=0), make_fn("a"))
        p.add_stage("b", create_cpu_stream(), StreamPlacement.cpu(rank=1), make_fn("b"), ["a"])
        p.add_stage("c", create_cpu_stream(), StreamPlacement.cpu(rank=2), make_fn("c"), ["b"])
        p.add_stage("d", create_cpu_stream(), StreamPlacement.npu(device_id=0), make_fn("d"), ["a"])
        p.add_stage("e", create_cpu_stream(), StreamPlacement.cuda(device_id=0), make_fn("e"))

        r = p.execute()

        assert r.success
        assert len(r.stage_results) == 5
        assert set(sr.stage_name for sr in r.stage_results) == {"a", "b", "c", "d", "e"}
        # Ordering constraints
        assert log.index("a") < log.index("b")
        assert log.index("b") < log.index("c")
        assert log.index("a") < log.index("d")

    def test_stage_failure_captured(self) -> None:
        """Failed stage is captured in StageResult without crashing pipeline."""
        p = MultiStreamPipeline(trace_id="fail-test")

        def bad_fn() -> None:
            raise ValueError("intentional error")

        p.add_stage("ok", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        p.add_stage("bad", create_cpu_stream(), StreamPlacement.cpu(), bad_fn)

        r = p.execute()

        assert not r.success  # At least one failure
        bad_result = next(sr for sr in r.stage_results if sr.stage_name == "bad")
        assert not bad_result.success
        assert "intentional error" in (bad_result.error or "")

    def test_result_has_trace_id(self) -> None:
        p = MultiStreamPipeline(trace_id="trace-xyz")
        p.add_stage("s", create_cpu_stream(), StreamPlacement.cpu(), lambda: None)
        r = p.execute()
        assert r.trace_id == "trace-xyz"

    def test_result_total_elapsed_positive(self) -> None:
        p = MultiStreamPipeline()
        p.add_stage("s", create_cpu_stream(), StreamPlacement.cpu(), lambda: time.sleep(0.01))
        r = p.execute()
        assert r.total_elapsed_ms >= 10.0

    def test_placement_summary_in_result(self) -> None:
        p = MultiStreamPipeline(trace_id="placement-check")
        pa = StreamPlacement.cpu(rank=0)
        pb = StreamPlacement.cuda(device_id=1)
        p.add_stage("comm", create_cpu_stream(), pa, lambda: None)
        p.add_stage("compute", create_cpu_stream(), pb, lambda: None, ["comm"])
        r = p.execute()
        assert r.placement_summary["comm"] == pa.key
        assert r.placement_summary["compute"] == pb.key

    def test_stage_elapsed_ms_measured(self) -> None:
        p = MultiStreamPipeline()
        p.add_stage("timed", create_cpu_stream(), StreamPlacement.cpu(), lambda: time.sleep(0.03))
        r = p.execute()
        assert r.stage_results[0].elapsed_ms >= 30.0


# =============================================================================
# Integration: 3-stream comm-compute-decode scenario
# =============================================================================


class TestThreeStreamCommComputeDecode:
    """Integration test mimicking a realistic comm / compute / decode pipeline."""

    def test_comm_compute_decode_pipeline(self) -> None:
        """
        3-stage pipeline:
          comm   (CPU, rank=0) – simulates collective communication
          compute (CPU, rank=1) – simulates computation after comm finishes
          decode  (CPU, rank=0) – simulates decode that depends on compute
        """
        timeline: list[str] = []
        lock = __import__("threading").Lock()

        def comm_fn() -> None:
            time.sleep(0.02)  # 20ms "communication"
            with lock:
                timeline.append("comm_done")

        def compute_fn() -> None:
            time.sleep(0.01)  # 10ms "compute"
            with lock:
                timeline.append("compute_done")

        def decode_fn() -> None:
            with lock:
                timeline.append("decode_done")

        p = MultiStreamPipeline(trace_id="comm-compute-decode")

        p.add_stage(
            "comm",
            create_cpu_stream(trace_id="comm-compute-decode"),
            StreamPlacement.cpu(rank=0),
            comm_fn,
        )
        p.add_stage(
            "compute",
            create_cpu_stream(trace_id="comm-compute-decode"),
            StreamPlacement.cpu(rank=1),
            compute_fn,
            dependencies=["comm"],
        )
        p.add_stage(
            "decode",
            create_cpu_stream(trace_id="comm-compute-decode"),
            StreamPlacement.cpu(rank=0),
            decode_fn,
            dependencies=["compute"],
        )

        result = p.execute()

        assert result.success
        assert result.trace_id == "comm-compute-decode"
        assert len(result.stage_results) == 3

        # Timeline order
        assert timeline == ["comm_done", "compute_done", "decode_done"]

        # Each stage should have a measured elapsed time
        for sr in result.stage_results:
            assert isinstance(sr, StageResult)
            assert sr.elapsed_ms >= 0.0

        # Placement summary present for all stages
        assert set(result.placement_summary.keys()) == {"comm", "compute", "decode"}


# =============================================================================
# Public API imports
# =============================================================================


class TestPublicAPIImports:
    """Smoke-test that all Task1.8 symbols are importable from the top-level package."""

    def test_import_from_sagellm_comm(self) -> None:
        from sagellm_comm import (  # noqa: F401
            MultiStreamPipeline,
            PipelineResult,
            PipelineStage,
            StageResult,
            StreamPlacement,
        )

    def test_import_from_overlap(self) -> None:
        from sagellm_comm.overlap import (  # noqa: F401
            DeviceType,
            MultiStreamPipeline,
            PipelineStage,
            StageResult,
            StreamPlacement,
        )
