"""Tests for agent lifecycle state machine.

Phase 5, step 3: State transitions and invariant enforcement.
"""

import pytest
from returns.result import Success

from anima.app.session_service import DefaultSessionService, TalkAttachRequest
from anima.core.contracts import SessionTransportState
from anima.core.execution import ExecutionSessionState
from anima.core.schemas.models import AgentStateData
from anima.shell.state import ShellSessionProjectionPort
from tests.fixtures.execution_contract_resolution import (
    make_consumer_conformance_fixture,
)

from anima.core.lifecycle import (
    TRANSITION_TABLE,
    TERMINAL_STATES,
    AgentState,
    LifecycleEvent,
    next_state,
    is_terminal,
    can_transition,
    valid_events_for,
    killable_states,
)
from anima.core.schemas import InvalidInputError


# =============================================================================
# Test Transition Table
# =============================================================================


class TestTransitionTable:
    """Tests for TRANSITION_TABLE structure and invariants."""

    def test_table_has_17_entries(self) -> None:
        """TRANSITION_TABLE must have exactly 17 entries."""
        assert len(TRANSITION_TABLE) == 17

    @pytest.mark.parametrize(
        "current,event,expected",
        [
            # From pending (1 transition)
            (AgentState.pending, LifecycleEvent.start, AgentState.running),
            # From running (10 transitions)
            (AgentState.running, LifecycleEvent.finish, AgentState.completed),
            (AgentState.running, LifecycleEvent.fail, AgentState.failed),
            (AgentState.running, LifecycleEvent.suspend, AgentState.suspended),
            (
                AgentState.running,
                LifecycleEvent.request_approval,
                AgentState.awaiting_approval,
            ),
            (AgentState.running, LifecycleEvent.timeout, AgentState.failed),
            (
                AgentState.running,
                LifecycleEvent.budget_exhausted,
                AgentState.failed,
            ),
            (AgentState.running, LifecycleEvent.max_turns, AgentState.failed),
            (AgentState.running, LifecycleEvent.crash, AgentState.failed),
            (AgentState.running, LifecycleEvent.kill, AgentState.killed),
            (AgentState.running, LifecycleEvent.auto_complete, AgentState.completed),
            # From suspended (2 transitions)
            (AgentState.suspended, LifecycleEvent.resume, AgentState.running),
            (AgentState.suspended, LifecycleEvent.kill, AgentState.killed),
            # From awaiting_approval (3 transitions)
            (
                AgentState.awaiting_approval,
                LifecycleEvent.approve,
                AgentState.running,
            ),
            (AgentState.awaiting_approval, LifecycleEvent.reject, AgentState.running),
            (AgentState.awaiting_approval, LifecycleEvent.kill, AgentState.killed),
        ],
    )
    def test_all_valid_transitions(
        self,
        current: AgentState,
        event: LifecycleEvent,
        expected: AgentState,
    ) -> None:
        """All 16 transitions produce expected next states."""
        assert TRANSITION_TABLE[(current, event)] == expected

    def test_terminal_states_have_no_outgoing(self) -> None:
        """Terminal states (completed, failed, killed) have no outgoing transitions."""
        terminal_states = {AgentState.completed, AgentState.failed, AgentState.killed}
        for state, _event in TRANSITION_TABLE:
            assert state not in terminal_states, (
                f"Terminal state {state} should have no outgoing transitions"
            )

    def test_pending_has_start_and_kill(self) -> None:
        """Pending state has 'start' and 'kill' as valid events."""
        valid_from_pending = valid_events_for(AgentState.pending)
        assert valid_from_pending == frozenset(
            {LifecycleEvent.start, LifecycleEvent.kill}
        )


# =============================================================================
# Test next_state Function
# =============================================================================


class TestNextState:
    """Tests for next_state transition function."""

    def test_valid_transition(self) -> None:
        """Valid transition returns expected state."""
        result = next_state(AgentState.pending, LifecycleEvent.start)
        assert result == AgentState.running

    def test_invalid_transition_raises(self) -> None:
        """Invalid transition raises InvalidInputError."""
        with pytest.raises(InvalidInputError) as exc_info:
            next_state(AgentState.completed, LifecycleEvent.start)
        assert "Invalid transition" in str(exc_info.value)

        with pytest.raises(InvalidInputError):
            next_state(AgentState.failed, LifecycleEvent.start)

        with pytest.raises(InvalidInputError):
            next_state(AgentState.killed, LifecycleEvent.start)

        # Invalid: no suspend from pending
        with pytest.raises(InvalidInputError):
            next_state(AgentState.pending, LifecycleEvent.suspend)

        # Invalid: no finish from pending
        with pytest.raises(InvalidInputError):
            next_state(AgentState.pending, LifecycleEvent.finish)

    @pytest.mark.parametrize(
        "state",
        [
            AgentState.running,
            AgentState.suspended,
            AgentState.awaiting_approval,
        ],
    )
    def test_kill_from_all_killable_states(self, state: AgentState) -> None:
        """Kill event transitions all killable states to killed."""
        result = next_state(state, LifecycleEvent.kill)
        assert result == AgentState.killed

    def test_kill_from_pending_valid(self) -> None:
        """Kill from pending is a valid transition (pending agents can be killed before start)."""
        assert (AgentState.pending, LifecycleEvent.kill) in TRANSITION_TABLE
        assert (
            TRANSITION_TABLE[(AgentState.pending, LifecycleEvent.kill)]
            == AgentState.killed
        )


# =============================================================================
# Test Helper Functions
# =============================================================================


class TestHelpers:
    """Tests for lifecycle helper functions."""

    # -------------------------------------------------------------------------
    # is_terminal tests
    # -------------------------------------------------------------------------

    @pytest.mark.parametrize(
        "state",
        [
            AgentState.completed,
            AgentState.failed,
            AgentState.killed,
        ],
    )
    def test_is_terminal_true(self, state: AgentState) -> None:
        """Terminal states return True for is_terminal."""
        assert is_terminal(state) is True
        assert state in TERMINAL_STATES

    @pytest.mark.parametrize(
        "state",
        [
            AgentState.pending,
            AgentState.running,
            AgentState.suspended,
            AgentState.awaiting_approval,
        ],
    )
    def test_is_terminal_false(self, state: AgentState) -> None:
        """Non-terminal states return False for is_terminal."""
        assert is_terminal(state) is False
        assert state not in TERMINAL_STATES

    # -------------------------------------------------------------------------
    # can_transition tests
    # -------------------------------------------------------------------------

    def test_can_transition_true(self) -> None:
        """Valid transitions return True."""
        assert can_transition(AgentState.pending, LifecycleEvent.start) is True
        assert can_transition(AgentState.running, LifecycleEvent.finish) is True
        assert can_transition(AgentState.running, LifecycleEvent.kill) is True
        assert can_transition(AgentState.suspended, LifecycleEvent.resume) is True
        assert (
            can_transition(AgentState.awaiting_approval, LifecycleEvent.approve) is True
        )

    def test_can_transition_false(self) -> None:
        """Invalid transitions return False."""
        # Terminal states have no valid transitions
        assert can_transition(AgentState.completed, LifecycleEvent.start) is False
        assert can_transition(AgentState.failed, LifecycleEvent.start) is False
        assert can_transition(AgentState.killed, LifecycleEvent.start) is False

        # No edge for this transition
        assert can_transition(AgentState.pending, LifecycleEvent.suspend) is False
        assert can_transition(AgentState.running, LifecycleEvent.start) is False

    # -------------------------------------------------------------------------
    # valid_events_for tests
    # -------------------------------------------------------------------------

    def test_valid_events_for_running(self) -> None:
        """Running state has 10 valid events."""
        events = valid_events_for(AgentState.running)
        expected_events = frozenset(
            {
                LifecycleEvent.finish,
                LifecycleEvent.fail,
                LifecycleEvent.suspend,
                LifecycleEvent.request_approval,
                LifecycleEvent.timeout,
                LifecycleEvent.budget_exhausted,
                LifecycleEvent.max_turns,
                LifecycleEvent.crash,
                LifecycleEvent.kill,
                LifecycleEvent.auto_complete,
            }
        )
        assert events == expected_events
        assert len(events) == 10

    def test_valid_events_for_suspended(self) -> None:
        """Suspended state has 2 valid events."""
        events = valid_events_for(AgentState.suspended)
        assert events == frozenset({LifecycleEvent.resume, LifecycleEvent.kill})
        assert len(events) == 2

    def test_valid_events_for_awaiting_approval(self) -> None:
        """Awaiting approval state has 3 valid events."""
        events = valid_events_for(AgentState.awaiting_approval)
        assert events == frozenset(
            {
                LifecycleEvent.approve,
                LifecycleEvent.reject,
                LifecycleEvent.kill,
            }
        )
        assert len(events) == 3

    @pytest.mark.parametrize(
        "state",
        [
            AgentState.completed,
            AgentState.failed,
            AgentState.killed,
        ],
    )
    def test_valid_events_for_terminal_empty(self, state: AgentState) -> None:
        """Terminal states have no valid events."""
        events = valid_events_for(state)
        assert events == frozenset()

    # -------------------------------------------------------------------------
    # killable_states tests
    # -------------------------------------------------------------------------

    def test_killable_states(self) -> None:
        """killable_states returns correct set of states."""
        states = killable_states()
        expected = frozenset(
            {
                AgentState.running,
                AgentState.suspended,
                AgentState.awaiting_approval,
                AgentState.pending,
            }
        )
        assert states == expected
        assert len(states) == 4

    def test_killable_states_includes_pending(self) -> None:
        """Pending is killable (no SIGTERM needed since pid is None)."""
        states = killable_states()
        assert AgentState.pending in states

    def test_killable_states_not_includes_terminal(self) -> None:
        """Terminal states are not killable."""
        states = killable_states()
        assert AgentState.completed not in states
        assert AgentState.failed not in states
        assert AgentState.killed not in states

    def test_killable_states_all_transition_to_killed(self) -> None:
        """All killable states transition to killed on kill event."""
        for state in killable_states():
            assert can_transition(state, LifecycleEvent.kill) is True
            assert next_state(state, LifecycleEvent.kill) == AgentState.killed


# =============================================================================
# Additional Coverage Tests
# =============================================================================


class TestLifecycleCoverage:
    """Additional tests for edge cases and coverage."""

    def test_transition_table_keys_are_valid(self) -> None:
        """All TRANSITION_TABLE keys are valid (state, event) pairs."""
        for state, event in TRANSITION_TABLE:
            assert isinstance(state, AgentState)
            assert isinstance(event, LifecycleEvent)

    def test_transition_table_values_are_valid(self) -> None:
        """All TRANSITION_TABLE values are valid states."""
        for next_st in TRANSITION_TABLE.values():
            assert isinstance(next_st, AgentState)

    def test_all_events_used(self) -> None:
        """All LifecycleEvent values appear at least once in TRANSITION_TABLE."""
        used_events = {event for (_, event) in TRANSITION_TABLE.keys()}
        all_events = set(LifecycleEvent)
        assert used_events == all_events

    def test_all_non_terminal_states_have_outgoing(self) -> None:
        """All non-terminal states have at least one outgoing transition."""
        non_terminal = {
            AgentState.pending,
            AgentState.running,
            AgentState.suspended,
            AgentState.awaiting_approval,
        }
        states_with_outgoing = {state for (state, _) in TRANSITION_TABLE}
        assert non_terminal == states_with_outgoing

    def test_terminal_states_constant_matches_enum(self) -> None:
        """TERMINAL_STATES matches the terminal states from AgentState enum."""
        expected_terminal = {
            AgentState.completed,
            AgentState.failed,
            AgentState.killed,
        }
        assert TERMINAL_STATES == expected_terminal


class TestLifecycleSupplementalInvariants:
    """Supplemental invariant checks for lifecycle transition coherence.

    Source: vectl step supplemental-tests-lifecycle-property guidance requests
    concise, high-ROI lifecycle invariant/property checks.
    """

    def test_transition_helpers_are_coherent_for_all_state_event_pairs(self) -> None:
        """For every (state, event), helpers agree on transition validity."""
        for state in AgentState:
            valid_events = valid_events_for(state)
            for event in LifecycleEvent:
                pair = (state, event)
                is_valid = pair in TRANSITION_TABLE

                assert can_transition(state, event) is is_valid
                assert (event in valid_events) is is_valid

                if is_valid:
                    assert next_state(state, event) == TRANSITION_TABLE[pair]
                else:
                    with pytest.raises(InvalidInputError):
                        next_state(state, event)

    def test_killability_matches_kill_event_reachability(self) -> None:
        """A state is killable iff kill is in its valid event set."""
        killable = killable_states()
        for state in AgentState:
            assert (state in killable) is (
                LifecycleEvent.kill in valid_events_for(state)
            )


class TestSessionAttachTerminalRegression:
    """Regression coverage for terminal-state-first attach semantics."""

    @pytest.mark.parametrize(
        "terminal_state",
        [AgentState.completed, AgentState.failed, AgentState.killed],
    )
    def test_terminal_state_preferred_when_socket_missing(
        self,
        terminal_state: AgentState,
    ) -> None:
        fixture = make_consumer_conformance_fixture(instance_id="inst_terminal_attach")
        mutable_state = AgentStateData(
            instance_id="inst_terminal_attach",
            trace_id="tr_terminal_attach",
            persona_name="terminal",
            state=terminal_state,
            mcp_socket="/tmp/does-not-exist.sock",
        )
        session_state = ExecutionSessionState(
            instance_id="inst_terminal_attach",
            transport=SessionTransportState(
                attach_mode="interactive",
                transport_connected=False,
                socket_path="/tmp/does-not-exist.sock",
            ),
        )

        class _Port(ShellSessionProjectionPort):
            def read_state(self, instance_id: str):
                _ = instance_id
                return Success(mutable_state)

            def load_contract(self, instance_id: str):
                _ = instance_id
                return Success(fixture.contract)

            def load_session_state(self, instance_id: str):
                _ = instance_id
                return Success(session_state)

        service = DefaultSessionService(projection_port=_Port())
        result = service.project_talk_attach(
            TalkAttachRequest(instance_id="inst_terminal_attach")
        )
        assert isinstance(result, Success)
        projection = result.unwrap()
        assert projection.presentation.kind == "terminal_state"
        assert projection.presentation.terminal_state == terminal_state
