"""
Additional unit tests for FSM execution functionality to improve coverage.
"""

import asyncio
from unittest.mock import MagicMock

import pytest

from fast_dag import FSM, ExecutionError, FSMContext, FSMReturn


class TestFSMBasicExecution:
    """Test basic FSM execution functionality"""

    def test_simple_fsm_execution(self):
        """Test simple FSM execution with state transitions"""
        fsm = FSM("simple_fsm")

        @fsm.state(initial=True)
        def start(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="process", value="started")

        @fsm.state
        def process(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="end", value="processed")

        @fsm.state(terminal=True)
        def end(context: FSMContext) -> FSMReturn:
            return FSMReturn(stop=True, value="done")

        result = fsm.run()
        assert result == "done"
        assert fsm.current_state == "end"

    def test_fsm_with_max_steps(self):
        """Test FSM execution with max_steps limit"""
        fsm = FSM("max_steps_fsm", max_cycles=5)

        @fsm.state(initial=True)
        def loop(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="loop", value="looping")

        with pytest.raises(ExecutionError) as exc_info:
            fsm.run()
        assert "Maximum cycles" in str(exc_info.value)

    def test_fsm_with_conditional_transitions(self):
        """Test FSM with conditional state transitions"""
        fsm = FSM("conditional_fsm")

        @fsm.state(initial=True)
        def check(context: FSMContext) -> FSMReturn:
            value = context.metadata.get("value", 0)
            if value > 0:
                return FSMReturn(next_state="positive_handler", value="positive")
            else:
                return FSMReturn(next_state="negative_handler", value="non_positive")

        @fsm.state(terminal=True)
        def positive_handler(context: FSMContext) -> FSMReturn:
            return FSMReturn(stop=True, value="handled_positive")

        @fsm.state(terminal=True)
        def negative_handler(context: FSMContext) -> FSMReturn:
            return FSMReturn(stop=True, value="handled_negative")

        # Test positive path
        context = FSMContext()
        context.metadata["value"] = 5
        result = fsm.run(context=context)
        assert result == "handled_positive"

        # Test negative path
        fsm.reset()
        context = FSMContext()
        context.metadata["value"] = -5
        result = fsm.run(context=context)
        assert result == "handled_negative"


class TestFSMAsyncExecution:
    """Test async FSM execution"""

    @pytest.mark.asyncio
    async def test_async_fsm_execution(self):
        """Test async FSM execution"""
        fsm = FSM("async_fsm")

        @fsm.state(initial=True)
        async def async_start(context: FSMContext) -> FSMReturn:
            await asyncio.sleep(0.01)
            return FSMReturn(next_state="async_process", value="async_started")

        @fsm.state(terminal=True)
        async def async_process(context: FSMContext) -> FSMReturn:
            await asyncio.sleep(0.01)
            return FSMReturn(stop=True, value="async_processed")

        result = await fsm.run_async()
        assert result == "async_processed"

    @pytest.mark.asyncio
    async def test_async_with_context(self):
        """Test async FSM with custom context"""
        fsm = FSM("async_context_fsm")

        @fsm.state(initial=True)
        async def node1(context: FSMContext) -> FSMReturn:
            context.metadata["value"] = 100
            return FSMReturn(next_state="node2", value=100)

        @fsm.state(terminal=True)
        async def node2(context: FSMContext) -> FSMReturn:
            prev_value = context.metadata.get("value", 0)
            return FSMReturn(stop=True, value=prev_value * 2)

        context = FSMContext()
        result = await fsm.run_async(context=context)
        assert result == 200


class TestFSMStepExecution:
    """Test step-by-step FSM execution"""

    def test_step_execution(self):
        """Test executing FSM step by step"""
        fsm = FSM("step_fsm")

        @fsm.state(initial=True)
        def step1(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="step2", value="step1_done")

        @fsm.state
        def step2(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="step3", value="step2_done")

        @fsm.state(terminal=True)
        def step3(context: FSMContext) -> FSMReturn:
            return FSMReturn(stop=True, value="all_done")

        context = FSMContext()

        # Execute step by step
        context, result1 = fsm.step(context)
        assert result1 == "step1_done"
        assert fsm.current_state == "step2"

        context, result2 = fsm.step(context)
        assert result2 == "step2_done"
        assert fsm.current_state == "step3"

        context, result3 = fsm.step(context)
        assert result3 == "all_done"
        assert fsm.is_terminal_state(fsm.current_state)

        # No more steps
        context, result4 = fsm.step(context)
        assert result4 is None


class TestFSMCycles:
    """Test FSM cycle handling"""

    def test_fsm_with_cycles(self):
        """Test FSM properly handles cycles"""
        fsm = FSM("cycle_fsm")

        @fsm.state(initial=True)
        def counter_node(context: FSMContext) -> FSMReturn:
            count = context.metadata.get("count", 0)
            count += 1
            context.metadata["count"] = count

            if count < 3:
                return FSMReturn(next_state="counter_node", value=count)
            else:
                return FSMReturn(stop=True, value="done")

        result = fsm.run()
        assert result == "done"
        assert fsm.context.metadata["count"] == 3

    def test_fsm_cycle_detection_in_step(self):
        """Test FSM detects cycles in step mode"""
        fsm = FSM("cycle_detection")

        @fsm.state(initial=True)
        def a(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="b", value="a")

        @fsm.state
        def b(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="a", value="b")

        context = FSMContext()

        # Should handle cycles gracefully in FSM
        context, result = fsm.step(context)
        assert result == "a"

        context, result = fsm.step(context)
        assert result == "b"

        # Can continue cycling
        context, result = fsm.step(context)
        assert result == "a"


class TestFSMErrorHandling:
    """Test FSM error handling"""

    def test_fsm_with_missing_next_node(self):
        """Test FSM handles missing next node"""
        fsm = FSM("missing_node_fsm")

        @fsm.state(initial=True)
        def start(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="missing", value="started")

        with pytest.raises(ExecutionError) as exc_info:
            fsm.run()
        assert "not found" in str(exc_info.value)

    def test_fsm_with_invalid_return(self):
        """Test FSM handles invalid return values"""
        fsm = FSM("invalid_return_fsm")

        @fsm.state(initial=True)
        def bad_node(context: FSMContext) -> str:  # Returns string instead of FSMReturn
            return "not_an_fsm_return"  # type: ignore

        # FSM might handle this differently - let's just check it doesn't crash
        from contextlib import suppress

        with suppress(ExecutionError, TypeError, AttributeError):
            fsm.run()
            # If it doesn't raise, that's okay for FSM

    def test_fsm_error_strategy_skip(self):
        """Test FSM with error during execution"""
        fsm = FSM("error_skip_fsm")

        @fsm.state(initial=True)
        def start(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="error_node", value="start")

        @fsm.state
        def error_node(context: FSMContext) -> FSMReturn:
            raise ValueError("Test error")

        # With default strategy, should raise
        with pytest.raises((ExecutionError, ValueError)):
            fsm.run()


class TestFSMSpecialFeatures:
    """Test FSM special features"""

    def test_fsm_dry_run(self):
        """Test FSM dry run functionality"""
        fsm = FSM("dry_run_fsm")

        @fsm.state(initial=True)
        def start(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="middle", value="start")

        @fsm.state
        def middle(context: FSMContext) -> FSMReturn:
            return FSMReturn(next_state="end", value="middle")

        @fsm.state(terminal=True)
        def end(context: FSMContext) -> FSMReturn:
            return FSMReturn(stop=True, value="end")

        # Check if FSM has dry_run method
        if hasattr(fsm, "dry_run"):
            plan = fsm.dry_run()
            # FSM dry_run returns different structure than DAG
            assert "node_details" in plan or "nodes" in plan
            assert len(plan.get("node_details", plan.get("nodes", {}))) == 3
        else:
            # If not implemented, just validate the FSM
            errors = fsm.validate()
            assert len(errors) == 0

    def test_fsm_get_metrics(self):
        """Test getting FSM execution metrics"""
        fsm = FSM("metrics_fsm")

        @fsm.state(initial=True)
        def timed_node(context: FSMContext) -> FSMReturn:
            import time

            time.sleep(0.01)
            return FSMReturn(stop=True, value="timed")

        result = fsm.run()
        assert result == "timed"

        # Check if FSM has get_metrics method
        if hasattr(fsm, "get_metrics"):
            metrics = fsm.get_metrics()
            assert isinstance(metrics, dict)
        else:
            # Just check that execution worked
            assert fsm.current_state is not None

    def test_fsm_reset(self):
        """Test FSM reset functionality"""
        fsm = FSM("reset_fsm")

        @fsm.state(initial=True)
        def counter(context: FSMContext) -> FSMReturn:
            count = context.metadata.get("count", 0) + 1
            context.metadata["count"] = count
            return FSMReturn(stop=True, value=count)

        result1 = fsm.run()
        assert result1 == 1

        # Reset and run again
        fsm.reset()
        result2 = fsm.run()
        assert result2 == 1  # Should start fresh

    def test_fsm_validation_errors(self):
        """Test FSM validation error handling"""
        fsm = FSM("invalid_fsm")

        # Create nodes but don't define them properly
        fsm.nodes["undefined"] = MagicMock()

        # This should trigger validation errors
        errors = fsm.validate()
        assert len(errors) > 0
