"""Tests for Pipeline Saga compensation pattern"""

import asyncio
from unittest.mock import AsyncMock

import pytest

from agentkit.orchestrator.compensation import (
    CompletedStep,
    CompensationResult,
    SagaOrchestrator,
)
from agentkit.orchestrator.pipeline_engine import PipelineEngine
from agentkit.orchestrator.pipeline_schema import (
    Pipeline,
    PipelineStage,
    StageStatus,
)
from agentkit.orchestrator.retry import StepRetryPolicy
from agentkit.skills.geo_pipeline import (
    GEO_PIPELINE_COMPENSATIONS,
    PipelineStep,
    create_geo_pipeline_steps,
)


class TestCompletedStep:
    """CompletedStep dataclass tests"""

    def test_creation_with_compensation(self):
        step = CompletedStep(
            step_name="optimize",
            result={"changes": 5},
            compensate_action="revert_optimization",
        )
        assert step.step_name == "optimize"
        assert step.result == {"changes": 5}
        assert step.compensate_action == "revert_optimization"

    def test_creation_without_compensation(self):
        step = CompletedStep(step_name="detect", result={"found": 3})
        assert step.compensate_action is None


class TestCompensationResult:
    """CompensationResult dataclass tests"""

    def test_success_result(self):
        result = CompensationResult(step_name="optimize", success=True)
        assert result.step_name == "optimize"
        assert result.success is True
        assert result.error is None

    def test_failure_result(self):
        result = CompensationResult(
            step_name="optimize", success=False, error="rollback failed"
        )
        assert result.success is False
        assert result.error == "rollback failed"


class TestSagaOrchestrator:
    """SagaOrchestrator tests"""

    def test_record_completed(self):
        saga = SagaOrchestrator()
        saga.record_completed("step1", {"data": 1}, "compensate_1")
        saga.record_completed("step2", {"data": 2})

        steps = saga.completed_steps
        assert len(steps) == 2
        assert steps[0].step_name == "step1"
        assert steps[0].compensate_action == "compensate_1"
        assert steps[1].step_name == "step2"
        assert steps[1].compensate_action is None

    @pytest.mark.asyncio
    async def test_compensate_lifo_order(self):
        """Compensation should execute in LIFO (reverse) order"""
        execution_order = []

        async def mock_execute_skill(skill_name: str, input_data):
            execution_order.append(skill_name)

        saga = SagaOrchestrator(execute_skill_func=mock_execute_skill)
        saga.record_completed("step1", {"data": 1}, "compensate_1")
        saga.record_completed("step2", {"data": 2}, "compensate_2")
        saga.record_completed("step3", {"data": 3}, "compensate_3")

        results = await saga.compensate()

        # LIFO order: step3, step2, step1
        assert execution_order == ["compensate_3", "compensate_2", "compensate_1"]
        assert len(results) == 3
        assert all(r.success for r in results)

    @pytest.mark.asyncio
    async def test_skip_steps_with_no_compensate_action(self):
        """Steps with no compensate_action should be skipped"""
        saga = SagaOrchestrator()
        saga.record_completed("read_only", {"data": 1})  # No compensation
        saga.record_completed("write_op", {"data": 2}, "rollback_write")

        results = await saga.compensate()

        assert len(results) == 2
        # write_op is compensated first (LIFO), then read_only
        assert results[0].step_name == "write_op"
        assert results[0].success is True
        assert results[1].step_name == "read_only"
        assert results[1].success is True
        assert results[1].error == "no_compensation_needed"

    @pytest.mark.asyncio
    async def test_compensation_failure_doesnt_interrupt_others(self):
        """If one compensation fails, others should still execute"""
        execution_order = []

        async def mock_execute_skill(skill_name: str, input_data):
            execution_order.append(skill_name)
            if skill_name == "compensate_2":
                raise RuntimeError("rollback failed")

        saga = SagaOrchestrator(execute_skill_func=mock_execute_skill)
        saga.record_completed("step1", {"data": 1}, "compensate_1")
        saga.record_completed("step2", {"data": 2}, "compensate_2")
        saga.record_completed("step3", {"data": 3}, "compensate_3")

        results = await saga.compensate()

        # All compensations should be attempted (LIFO: step3, step2, step1)
        assert execution_order == ["compensate_3", "compensate_2", "compensate_1"]
        assert len(results) == 3

        # step3 succeeds
        assert results[0].success is True
        # step2 fails
        assert results[1].success is False
        assert results[1].error == "rollback failed"
        # step1 still succeeds
        assert results[2].success is True

    @pytest.mark.asyncio
    async def test_compensate_with_no_execute_skill_func(self):
        """Without execute_skill_func, compensation succeeds but does nothing"""
        saga = SagaOrchestrator()
        saga.record_completed("step1", {"data": 1}, "compensate_1")

        results = await saga.compensate()

        assert len(results) == 1
        assert results[0].success is True

    def test_clear(self):
        saga = SagaOrchestrator()
        saga.record_completed("step1", {"data": 1})
        saga.record_completed("step2", {"data": 2})
        assert len(saga.completed_steps) == 2

        saga.clear()
        assert len(saga.completed_steps) == 0

    def test_completed_steps_returns_copy(self):
        saga = SagaOrchestrator()
        saga.record_completed("step1", {"data": 1})

        steps = saga.completed_steps
        steps.clear()  # Mutate the copy

        assert len(saga.completed_steps) == 1  # Original unchanged


class TestPipelineIntegration:
    """Pipeline engine integration with retry and compensation"""

    @pytest.mark.asyncio
    async def test_step_failure_triggers_compensation(self):
        """When a step fails, Saga compensation should be triggered for completed steps"""
        engine = PipelineEngine()

        pipeline = Pipeline(
            name="test_compensation",
            version="1.0",
            description="Test compensation",
            stages=[
                PipelineStage(
                    name="step1",
                    agent="agent_a",
                    action="do_a",
                    compensate="undo_a",
                ),
                PipelineStage(
                    name="step2",
                    agent="agent_b",
                    action="do_b",
                    depends_on=["step1"],
                ),
            ],
        )

        # Dry-run mode (no dispatcher) — all steps succeed
        result = await engine.execute(pipeline)
        assert result.status == StageStatus.COMPLETED

    @pytest.mark.asyncio
    async def test_continue_on_failure(self):
        """Steps with continue_on_failure should not abort the pipeline"""
        engine = PipelineEngine()

        pipeline = Pipeline(
            name="test_continue",
            version="1.0",
            description="Test continue_on_failure",
            stages=[
                PipelineStage(
                    name="step1",
                    agent="agent_a",
                    action="do_a",
                    continue_on_failure=True,
                ),
                PipelineStage(
                    name="step2",
                    agent="agent_b",
                    action="do_b",
                    depends_on=["step1"],
                ),
            ],
        )

        # Dry-run mode — all steps succeed
        result = await engine.execute(pipeline)
        assert result.status == StageStatus.COMPLETED

    @pytest.mark.asyncio
    async def test_pipeline_with_retry_policy(self):
        """PipelineStage can have a retry_policy configured"""
        retry = StepRetryPolicy(max_attempts=5, base_delay=0.01, jitter=False)

        stage = PipelineStage(
            name="retry_step",
            agent="agent_a",
            action="do_a",
            retry_policy=retry,
        )
        assert stage.retry_policy is not None
        assert stage.retry_policy.max_attempts == 5

    @pytest.mark.asyncio
    async def test_pipeline_with_compensate(self):
        """PipelineStage can have a compensate action configured"""
        stage = PipelineStage(
            name="optimizable_step",
            agent="agent_a",
            action="do_a",
            compensate="undo_a",
        )
        assert stage.compensate == "undo_a"

    @pytest.mark.asyncio
    async def test_pipeline_without_compensate(self):
        """PipelineStage without compensate defaults to None"""
        stage = PipelineStage(
            name="readonly_step",
            agent="agent_a",
            action="do_a",
        )
        assert stage.compensate is None


class TestGEOPipelineCompensations:
    """GEO Pipeline compensation definitions"""

    def test_compensation_definitions(self):
        """Verify GEO pipeline compensation definitions"""
        assert GEO_PIPELINE_COMPENSATIONS["detect"] is None
        assert GEO_PIPELINE_COMPENSATIONS["analyze_competitor"] is None
        assert GEO_PIPELINE_COMPENSATIONS["optimize"] == "revert_optimization"
        assert GEO_PIPELINE_COMPENSATIONS["schema"] is None
        assert GEO_PIPELINE_COMPENSATIONS["monitor"] is None

    def test_create_geo_pipeline_steps(self):
        """Verify GEO pipeline steps are created with compensation"""
        steps = create_geo_pipeline_steps()
        assert len(steps) == 5

        step_names = [s.name for s in steps]
        assert step_names == [
            "detect",
            "analyze_competitor",
            "optimize",
            "schema",
            "monitor",
        ]

        # Check compensation assignments
        step_map = {s.name: s for s in steps}
        assert step_map["detect"].compensate is None
        assert step_map["analyze_competitor"].compensate is None
        assert step_map["optimize"].compensate == "revert_optimization"
        assert step_map["schema"].compensate is None
        assert step_map["monitor"].compensate is None

    def test_geo_pipeline_steps_dependencies(self):
        """Verify GEO pipeline step dependencies form a valid chain"""
        steps = create_geo_pipeline_steps()
        step_map = {s.name: s for s in steps}

        assert step_map["detect"].depends_on == []
        assert step_map["analyze_competitor"].depends_on == ["detect"]
        assert step_map["optimize"].depends_on == ["analyze_competitor"]
        assert step_map["schema"].depends_on == ["optimize"]
        assert step_map["monitor"].depends_on == ["schema"]