"""Unit tests for pipeline executor."""

import asyncio
import pytest
from typing import Any

from aperion_flow.definitions.pipeline import (
    LLMAssistanceConfig,
    LLMAssistanceMode,
    Pipeline,
    PipelineStage,
    PipelineStatus,
)
from aperion_flow.engine.context import InMemoryCheckpointStore, PipelineContext
from aperion_flow.engine.executor import ExecutionResult, PipelineExecutor


class TestPipelineExecutor:
    """Tests for PipelineExecutor."""

    @pytest.fixture
    def executor(self) -> PipelineExecutor:
        """Create an executor with in-memory store."""
        return PipelineExecutor(
            checkpoint_store=InMemoryCheckpointStore(),
        )

    @pytest.fixture
    def simple_pipeline(self) -> Pipeline:
        """Create a simple two-stage pipeline."""
        return Pipeline(
            name="simple_pipeline",
            stages=[
                PipelineStage(
                    id="s1",
                    name="stage1",
                    handler="tests.unit.test_executor.handler_success",
                ),
                PipelineStage(
                    id="s2",
                    name="stage2",
                    handler="tests.unit.test_executor.handler_success",
                    depends_on=["s1"],
                ),
            ],
        )

    @pytest.mark.asyncio
    async def test_execute_simple_pipeline(self, executor: PipelineExecutor) -> None:
        """Test executing a simple pipeline with registered handlers."""
        # Register handlers
        async def handler1(data: dict[str, Any]) -> dict[str, Any]:
            return {"stage1_output": "value1"}

        async def handler2(data: dict[str, Any]) -> dict[str, Any]:
            return {"stage2_output": "value2"}

        executor.register_handler("handler1", handler1)
        executor.register_handler("handler2", handler2)

        pipeline = Pipeline(
            name="test_pipeline",
            stages=[
                PipelineStage(id="s1", name="stage1", handler="handler1"),
                PipelineStage(id="s2", name="stage2", handler="handler2", depends_on=["s1"]),
            ],
        )

        result = await executor.execute(pipeline)

        assert result.status == PipelineStatus.COMPLETED
        assert result.stages_completed == 2
        assert result.stages_failed == 0
        assert "stage1_output" in result.data
        assert "stage2_output" in result.data

    @pytest.mark.asyncio
    async def test_execute_with_initial_data(self, executor: PipelineExecutor) -> None:
        """Test executing with initial data."""
        async def handler(data: dict[str, Any]) -> dict[str, Any]:
            return {"received": data.get("input_value")}

        executor.register_handler("handler", handler)

        pipeline = Pipeline(
            name="test_pipeline",
            stages=[PipelineStage(id="s1", name="stage1", handler="handler")],
        )

        result = await executor.execute(
            pipeline,
            initial_data={"input_value": "test_input"},
        )

        assert result.status == PipelineStatus.COMPLETED
        assert result.data.get("received") == "test_input"

    @pytest.mark.asyncio
    async def test_execute_parallel_stages(self, executor: PipelineExecutor) -> None:
        """Test executing stages in parallel."""
        execution_order: list[str] = []

        async def handler_a(data: dict[str, Any]) -> dict[str, Any]:
            execution_order.append("a_start")
            await asyncio.sleep(0.05)
            execution_order.append("a_end")
            return {"a": "done"}

        async def handler_b(data: dict[str, Any]) -> dict[str, Any]:
            execution_order.append("b_start")
            await asyncio.sleep(0.05)
            execution_order.append("b_end")
            return {"b": "done"}

        executor.register_handler("handler_a", handler_a)
        executor.register_handler("handler_b", handler_b)

        pipeline = Pipeline(
            name="parallel_pipeline",
            stages=[
                PipelineStage(id="s1", name="stage_a", handler="handler_a"),
                PipelineStage(id="s2", name="stage_b", handler="handler_b"),
            ],
            max_parallel=2,
        )

        result = await executor.execute(pipeline)

        assert result.status == PipelineStatus.COMPLETED
        # Both should start before either ends (parallel execution)
        assert "a_start" in execution_order[:2]
        assert "b_start" in execution_order[:2]

    @pytest.mark.asyncio
    async def test_execute_with_dependency(self, executor: PipelineExecutor) -> None:
        """Test that dependencies are respected."""
        execution_order: list[str] = []

        async def handler_first(data: dict[str, Any]) -> dict[str, Any]:
            execution_order.append("first")
            return {"first": True}

        async def handler_second(data: dict[str, Any]) -> dict[str, Any]:
            execution_order.append("second")
            assert data.get("first") is True  # Verify dependency data
            return {"second": True}

        executor.register_handler("handler_first", handler_first)
        executor.register_handler("handler_second", handler_second)

        pipeline = Pipeline(
            name="dependency_pipeline",
            stages=[
                PipelineStage(id="s1", name="first", handler="handler_first"),
                PipelineStage(id="s2", name="second", handler="handler_second", depends_on=["s1"]),
            ],
        )

        result = await executor.execute(pipeline)

        assert result.status == PipelineStatus.COMPLETED
        assert execution_order == ["first", "second"]

    @pytest.mark.asyncio
    async def test_execute_stage_failure(self, executor: PipelineExecutor) -> None:
        """Test handling stage failure."""
        async def failing_handler(data: dict[str, Any]) -> dict[str, Any]:
            raise ValueError("Stage failed!")

        executor.register_handler("failing", failing_handler)

        pipeline = Pipeline(
            name="failing_pipeline",
            stages=[PipelineStage(id="s1", name="failing_stage", handler="failing")],
        )

        result = await executor.execute(pipeline)

        assert result.status == PipelineStatus.FAILED
        assert result.stages_failed == 1
        assert result.error is not None
        assert "Stage failed!" in result.error

    @pytest.mark.asyncio
    async def test_execute_fail_fast(self, executor: PipelineExecutor) -> None:
        """Test fail_fast behavior."""
        executed: list[str] = []

        async def handler_fail(data: dict[str, Any]) -> dict[str, Any]:
            executed.append("fail")
            raise ValueError("Failure")

        async def handler_after(data: dict[str, Any]) -> dict[str, Any]:
            executed.append("after")
            return {}

        executor.register_handler("handler_fail", handler_fail)
        executor.register_handler("handler_after", handler_after)

        pipeline = Pipeline(
            name="fail_fast_pipeline",
            stages=[
                PipelineStage(id="s1", name="failing", handler="handler_fail"),
                PipelineStage(id="s2", name="after", handler="handler_after", depends_on=["s1"]),
            ],
            fail_fast=True,
        )

        result = await executor.execute(pipeline)

        assert result.status == PipelineStatus.FAILED
        assert "after" not in executed

    @pytest.mark.asyncio
    async def test_execute_skip_on_failure(self, executor: PipelineExecutor) -> None:
        """Test skip_on_failure behavior."""
        async def handler_fail(data: dict[str, Any]) -> dict[str, Any]:
            raise ValueError("Failure")

        async def handler_skip(data: dict[str, Any]) -> dict[str, Any]:
            return {"should_not_run": True}

        executor.register_handler("handler_fail", handler_fail)
        executor.register_handler("handler_skip", handler_skip)

        pipeline = Pipeline(
            name="skip_pipeline",
            stages=[
                PipelineStage(id="s1", name="failing", handler="handler_fail"),
                PipelineStage(
                    id="s2",
                    name="skippable",
                    handler="handler_skip",
                    depends_on=["s1"],
                    skip_on_failure=True,
                ),
            ],
            fail_fast=False,
        )

        result = await executor.execute(pipeline)

        assert result.stages_skipped == 1
        assert "should_not_run" not in result.data

    @pytest.mark.asyncio
    async def test_execute_with_timeout(self, executor: PipelineExecutor) -> None:
        """Test stage timeout handling."""
        async def slow_handler(data: dict[str, Any]) -> dict[str, Any]:
            await asyncio.sleep(10)  # Too slow
            return {"completed": True}

        executor.register_handler("slow", slow_handler)

        pipeline = Pipeline(
            name="timeout_pipeline",
            stages=[
                PipelineStage(
                    id="s1",
                    name="slow_stage",
                    handler="slow",
                    timeout=0.1,  # 100ms timeout
                ),
            ],
        )

        result = await executor.execute(pipeline)

        assert result.status == PipelineStatus.FAILED
        assert result.stages_failed == 1

    @pytest.mark.asyncio
    async def test_execute_with_correlation_id(self, executor: PipelineExecutor) -> None:
        """Test correlation ID is passed through."""
        async def handler(data: dict[str, Any]) -> dict[str, Any]:
            return {}

        executor.register_handler("handler", handler)

        pipeline = Pipeline(
            name="correlated_pipeline",
            stages=[PipelineStage(id="s1", name="stage1", handler="handler")],
        )

        result = await executor.execute(
            pipeline,
            correlation_id="test-correlation-123",
        )

        assert result.status == PipelineStatus.COMPLETED

    @pytest.mark.asyncio
    async def test_checkpointing(self, executor: PipelineExecutor) -> None:
        """Test that checkpoints are created."""
        async def handler(data: dict[str, Any]) -> dict[str, Any]:
            return {"output": "value"}

        executor.register_handler("handler", handler)

        pipeline = Pipeline(
            name="checkpointed_pipeline",
            stages=[PipelineStage(id="s1", name="stage1", handler="handler")],
        )

        result = await executor.execute(pipeline)

        # Verify checkpoints were saved
        store = executor._checkpoint_store
        checkpoints = await store.list_checkpoints(result.execution_id)
        assert len(checkpoints) > 0

    @pytest.mark.asyncio
    async def test_handler_not_found(self, executor: PipelineExecutor) -> None:
        """Test error when handler is not found."""
        pipeline = Pipeline(
            name="missing_handler_pipeline",
            stages=[
                PipelineStage(
                    id="s1",
                    name="stage1",
                    handler="nonexistent.module.handler",
                ),
            ],
        )

        result = await executor.execute(pipeline)

        assert result.status == PipelineStatus.FAILED
        assert "Cannot load handler" in (result.error or "")


class TestExecutionResult:
    """Tests for ExecutionResult."""

    def test_result_structure(self) -> None:
        """Test result has expected structure."""
        result = ExecutionResult(
            execution_id="e1",
            pipeline_id="p1",
            pipeline_name="test",
            status=PipelineStatus.COMPLETED,
            stage_count=3,
            stages_completed=3,
            stages_failed=0,
            stages_skipped=0,
        )

        assert result.execution_id == "e1"
        assert result.status == PipelineStatus.COMPLETED
        assert result.stage_count == 3

    def test_result_duration(self) -> None:
        """Test duration calculation."""
        from datetime import datetime, timezone

        start = datetime.now(timezone.utc)
        end = datetime.now(timezone.utc)

        result = ExecutionResult(
            execution_id="e1",
            pipeline_id="p1",
            pipeline_name="test",
            status=PipelineStatus.COMPLETED,
            started_at=start,
            completed_at=end,
        )

        assert result.duration_seconds is not None
        assert result.duration_seconds >= 0


class TestHandlerRegistration:
    """Tests for handler registration."""

    def test_register_handler(self) -> None:
        """Test registering a handler."""
        executor = PipelineExecutor()

        async def my_handler(data: dict[str, Any]) -> dict[str, Any]:
            return {"result": "success"}

        executor.register_handler("my_handler", my_handler)

        assert "my_handler" in executor._handler_registry

    @pytest.mark.asyncio
    async def test_registered_handler_used(self) -> None:
        """Test that registered handlers are used over dynamic import."""
        executor = PipelineExecutor()

        async def custom_handler(data: dict[str, Any]) -> dict[str, Any]:
            return {"source": "custom"}

        executor.register_handler("custom", custom_handler)

        pipeline = Pipeline(
            name="custom_handler_pipeline",
            stages=[PipelineStage(id="s1", name="stage1", handler="custom")],
        )

        result = await executor.execute(pipeline)

        assert result.status == PipelineStatus.COMPLETED
        assert result.data.get("source") == "custom"


# Helper handlers for tests
async def handler_success(data: dict[str, Any]) -> dict[str, Any]:
    """Simple successful handler."""
    return {"success": True}


async def handler_failure(data: dict[str, Any]) -> dict[str, Any]:
    """Handler that always fails."""
    raise ValueError("Handler failed")
