"""Agent lifecycle state machine (pure transitions, no I/O imports)."""

# @invar:allow file_size_warning: Lifecycle transition table intentionally co-located with event helpers for docs/INTERFACES.md H.1 parity

from enum import Enum

from anima.core.schemas import AgentState, InvalidInputError
from anima.core.schemas.enums import CompletionPolicy, LifecycleType

import deal


class LifecycleEvent(str, Enum):
    """Events that trigger state transitions in the agent lifecycle.

    Each event corresponds to an action or condition that may cause
    an agent to transition from one state to another.

    >>> LifecycleEvent.start.value
    'start'
    >>> LifecycleEvent.finish.value
    'finish'
    >>> LifecycleEvent.fail.value
    'fail'
    >>> LifecycleEvent.suspend.value
    'suspend'
    >>> LifecycleEvent.request_approval.value
    'request_approval'
    >>> LifecycleEvent.approve.value
    'approve'
    >>> LifecycleEvent.reject.value
    'reject'
    >>> LifecycleEvent.resume.value
    'resume'
    >>> LifecycleEvent.kill.value
    'kill'
    >>> LifecycleEvent.timeout.value
    'timeout'
    >>> LifecycleEvent.budget_exhausted.value
    'budget_exhausted'
    >>> LifecycleEvent.max_turns.value
    'max_turns'
    >>> LifecycleEvent.crash.value
    'crash'
    >>> LifecycleEvent.auto_complete.value
    'auto_complete'
    """

    start = "start"
    finish = "finish"
    fail = "fail"
    suspend = "suspend"
    request_approval = "request_approval"
    approve = "approve"
    reject = "reject"
    resume = "resume"
    kill = "kill"
    timeout = "timeout"
    budget_exhausted = "budget_exhausted"
    max_turns = "max_turns"
    crash = "crash"
    auto_complete = "auto_complete"


TRANSITION_TABLE: dict[tuple[AgentState, LifecycleEvent], AgentState] = {
    # From pending
    (AgentState.pending, LifecycleEvent.start): AgentState.running,
    (AgentState.pending, LifecycleEvent.kill): AgentState.killed,
    # From running
    (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
    (AgentState.suspended, LifecycleEvent.resume): AgentState.running,
    (AgentState.suspended, LifecycleEvent.kill): AgentState.killed,
    # From awaiting_approval
    (AgentState.awaiting_approval, LifecycleEvent.approve): AgentState.running,
    (AgentState.awaiting_approval, LifecycleEvent.reject): AgentState.running,
    (AgentState.awaiting_approval, LifecycleEvent.kill): AgentState.killed,
}
"""Mapping of (current_state, event) -> next_state for valid transitions.

>>> len(TRANSITION_TABLE)
17
>>> TRANSITION_TABLE[(AgentState.pending, LifecycleEvent.start)]
<AgentState.running: 'running'>
>>> TRANSITION_TABLE[(AgentState.running, LifecycleEvent.finish)]
<AgentState.completed: 'completed'>
"""


@deal.pre(lambda current, event: len(current.value) > 0 and len(event.value) > 0)
@deal.ensure(
    lambda current, event, result: TRANSITION_TABLE.get((current, event)) == result
)
@deal.raises(InvalidInputError)
def next_state(current: AgentState, event: LifecycleEvent) -> AgentState:
    """Return the next state for a valid transition.

    Looks up (current, event) in TRANSITION_TABLE and returns the next state.
    Raises InvalidInputError if the transition is not valid.

    Args:
        current: The current agent state.
        event: The lifecycle event that occurred.

    Returns:
        The next agent state.

    Raises:
        InvalidInputError: If (current, event) is not a valid transition.

    Valid transitions:
        >>> next_state(AgentState.pending, LifecycleEvent.start)
        <AgentState.running: 'running'>
        >>> next_state(AgentState.running, LifecycleEvent.finish)
        <AgentState.completed: 'completed'>
        >>> next_state(AgentState.running, LifecycleEvent.fail)
        <AgentState.failed: 'failed'>
        >>> next_state(AgentState.running, LifecycleEvent.suspend)
        <AgentState.suspended: 'suspended'>
        >>> next_state(AgentState.running, LifecycleEvent.request_approval)
        <AgentState.awaiting_approval: 'awaiting_approval'>
        >>> next_state(AgentState.running, LifecycleEvent.timeout)
        <AgentState.failed: 'failed'>
        >>> next_state(AgentState.running, LifecycleEvent.budget_exhausted)
        <AgentState.failed: 'failed'>
        >>> next_state(AgentState.running, LifecycleEvent.max_turns)
        <AgentState.failed: 'failed'>
        >>> next_state(AgentState.running, LifecycleEvent.crash)
        <AgentState.failed: 'failed'>
        >>> next_state(AgentState.running, LifecycleEvent.kill)
        <AgentState.killed: 'killed'>
        >>> next_state(AgentState.running, LifecycleEvent.auto_complete)
        <AgentState.completed: 'completed'>
        >>> next_state(AgentState.suspended, LifecycleEvent.resume)
        <AgentState.running: 'running'>
        >>> next_state(AgentState.suspended, LifecycleEvent.kill)
        <AgentState.killed: 'killed'>
        >>> next_state(AgentState.awaiting_approval, LifecycleEvent.approve)
        <AgentState.running: 'running'>
        >>> next_state(AgentState.awaiting_approval, LifecycleEvent.reject)
        <AgentState.running: 'running'>
        >>> next_state(AgentState.awaiting_approval, LifecycleEvent.kill)
        <AgentState.killed: 'killed'>

    Invalid transitions from terminal states:
        >>> next_state(AgentState.completed, LifecycleEvent.start)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        InvalidInputError: ...
        >>> next_state(AgentState.failed, LifecycleEvent.start)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        InvalidInputError: ...
        >>> next_state(AgentState.killed, LifecycleEvent.start)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        InvalidInputError: ...

    Invalid transition (no suspend from pending):
        >>> next_state(AgentState.pending, LifecycleEvent.suspend)  # doctest: +IGNORE_EXCEPTION_DETAIL
        Traceback (most recent call last):
        ...
        InvalidInputError: ...
    """
    key = (current, event)
    if key not in TRANSITION_TABLE:
        raise InvalidInputError(
            f"Invalid transition: cannot apply '{event.value}' in state '{current.value}'"
        )
    return TRANSITION_TABLE[key]


TERMINAL_STATES: frozenset[AgentState] = frozenset(
    {
        AgentState.completed,
        AgentState.failed,
        AgentState.killed,
    }
)
"""States from which no transitions are possible.

A terminal state has no outgoing edges in the state machine.
Once an agent enters a terminal state, it cannot transition to any other state.

>>> AgentState.completed in TERMINAL_STATES
True
>>> AgentState.failed in TERMINAL_STATES
True
>>> AgentState.killed in TERMINAL_STATES
True
>>> AgentState.running in TERMINAL_STATES
False
>>> AgentState.pending in TERMINAL_STATES
False
>>> AgentState.suspended in TERMINAL_STATES
False
>>> AgentState.awaiting_approval in TERMINAL_STATES
False
>>> len(TERMINAL_STATES)
3
"""


@deal.pre(lambda state: len(state.value) > 0)
@deal.ensure(lambda state, result: result == (state in TERMINAL_STATES))
def is_terminal(state: AgentState) -> bool:
    """Check if a state is terminal (no outgoing transitions).

    Terminal states are: completed, failed, killed.
    Once an agent enters a terminal state, it cannot transition further.

    Args:
        state: The agent state to check.

    Returns:
        True if the state is terminal, False otherwise.

    >>> is_terminal(AgentState.completed)
    True
    >>> is_terminal(AgentState.failed)
    True
    >>> is_terminal(AgentState.killed)
    True
    >>> is_terminal(AgentState.running)
    False
    >>> is_terminal(AgentState.pending)
    False
    >>> is_terminal(AgentState.suspended)
    False
    >>> is_terminal(AgentState.awaiting_approval)
    False
    """
    return state in TERMINAL_STATES


@deal.pre(lambda current, event: len(current.value) > 0 and len(event.value) > 0)
@deal.ensure(
    lambda current, event, result: result == ((current, event) in TRANSITION_TABLE)
)
def can_transition(current: AgentState, event: LifecycleEvent) -> bool:
    """Check if a transition from current state with given event is valid.

    Looks up (current, event) in TRANSITION_TABLE without raising an error.

    Args:
        current: The current agent state.
        event: The lifecycle event to check.

    Returns:
        True if the transition is valid, False otherwise.

    Valid transitions:
        >>> can_transition(AgentState.pending, LifecycleEvent.start)
        True
        >>> can_transition(AgentState.running, LifecycleEvent.finish)
        True
        >>> can_transition(AgentState.running, LifecycleEvent.kill)
        True
        >>> can_transition(AgentState.suspended, LifecycleEvent.resume)
        True
        >>> can_transition(AgentState.awaiting_approval, LifecycleEvent.approve)
        True

    Invalid transitions (terminal states):
        >>> can_transition(AgentState.completed, LifecycleEvent.start)
        False
        >>> can_transition(AgentState.failed, LifecycleEvent.start)
        False
        >>> can_transition(AgentState.killed, LifecycleEvent.start)
        False

    Invalid transitions (no edge):
        >>> can_transition(AgentState.pending, LifecycleEvent.suspend)
        False
        >>> can_transition(AgentState.running, LifecycleEvent.start)
        False
    """
    return (current, event) in TRANSITION_TABLE


@deal.pre(lambda state: len(state.value) > 0)
@deal.ensure(
    lambda state, result: (
        result
        == frozenset(
            event for (source_state, event) in TRANSITION_TABLE if source_state == state
        )
    )
)
def valid_events_for(state: AgentState) -> frozenset[LifecycleEvent]:
    """Return all valid events that can be applied from the given state.

    For terminal states, returns an empty frozenset.

    Args:
        state: The agent state to query.

    Returns:
        Frozenset of LifecycleEvents valid from this state.

    >>> valid_events_for(AgentState.pending) == frozenset({LifecycleEvent.start, LifecycleEvent.kill})
    True
    >>> valid_events_for(AgentState.completed)
    frozenset()
    >>> valid_events_for(AgentState.failed)
    frozenset()
    >>> valid_events_for(AgentState.killed)
    frozenset()

    Running state has many valid events:
        >>> events = valid_events_for(AgentState.running)
        >>> LifecycleEvent.finish in events
        True
        >>> LifecycleEvent.fail in events
        True
        >>> LifecycleEvent.suspend in events
        True
        >>> LifecycleEvent.kill in events
        True
        >>> LifecycleEvent.timeout in events
        True

    Suspended state:
        >>> events = valid_events_for(AgentState.suspended)
        >>> LifecycleEvent.resume in events
        True
        >>> LifecycleEvent.kill in events
        True
        >>> len(events)
        2

    Awaiting approval state:
        >>> events = valid_events_for(AgentState.awaiting_approval)
        >>> LifecycleEvent.approve in events
        True
        >>> LifecycleEvent.reject in events
        True
        >>> LifecycleEvent.kill in events
        True
        >>> len(events)
        3
    """
    return frozenset(event for (s, event) in TRANSITION_TABLE if s == state)


@deal.post(lambda result: isinstance(result, frozenset))
@deal.post(lambda result: all(isinstance(s, AgentState) for s in result))
@deal.post(lambda result: all(not is_terminal(s) for s in result))
def killable_states() -> frozenset[AgentState]:
    """Return all states from which a kill event can be triggered.

    These are states that have a valid transition with LifecycleEvent.kill.
    Terminal states are never killable (they cannot transition).
    Pending agents can also be killed (transition without SIGTERM since pid is None).

    Returns:
        Frozenset of AgentStates that can be killed.

    >>> states = killable_states()
    >>> AgentState.running in states
    True
    >>> AgentState.suspended in states
    True
    >>> AgentState.awaiting_approval in states
    True
    >>> AgentState.pending in states
    True
    >>> AgentState.completed in states
    False
    >>> AgentState.failed in states
    False
    >>> AgentState.killed in states
    False

    All killable states transition to killed:
        >>> for s in killable_states():
        ...     assert can_transition(s, LifecycleEvent.kill)
        ...     assert next_state(s, LifecycleEvent.kill) == AgentState.killed
    """
    return frozenset(
        state for (state, event) in TRANSITION_TABLE if event == LifecycleEvent.kill
    )


@deal.ensure(
    lambda requested_lifecycle=None, job_lifecycle=None, requested_completion_policy=None, job_completion_policy=None, result=None: (
        isinstance(result, tuple) and len(result) == 2
    )
)
@deal.post(lambda result: isinstance(result, tuple) and len(result) == 2)
def resolve_lifecycle_completion_policy(
    *,
    requested_lifecycle: LifecycleType | None,
    job_lifecycle: LifecycleType | None,
    requested_completion_policy: CompletionPolicy | None,
    job_completion_policy: CompletionPolicy | None,
) -> tuple[LifecycleType, CompletionPolicy]:
    """Resolve lifecycle and completion policy exactly once.

    Spec source:
    docs/proposals/execution-architecture-refactor.md#execution-contract-specific-design-rules/lifecycle-semantics

    Rules:
    - Explicit requested lifecycle always wins over job/default values.
    - Completion policy resolves separately and is never used to silently rewrite
      an explicit lifecycle.
    - If completion is unspecified, daemon defaults to explicit and ephemeral
      defaults to auto.

    >>> resolve_lifecycle_completion_policy(
    ...     requested_lifecycle=LifecycleType.daemon,
    ...     job_lifecycle=LifecycleType.ephemeral,
    ...     requested_completion_policy=None,
    ...     job_completion_policy=None,
    ... )
    (<LifecycleType.daemon: 'daemon'>, <CompletionPolicy.explicit: 'explicit'>)
    >>> resolve_lifecycle_completion_policy(
    ...     requested_lifecycle=None,
    ...     job_lifecycle=LifecycleType.ephemeral,
    ...     requested_completion_policy=None,
    ...     job_completion_policy=None,
    ... )
    (<LifecycleType.ephemeral: 'ephemeral'>, <CompletionPolicy.auto: 'auto'>)
    """
    lifecycle = requested_lifecycle or job_lifecycle or LifecycleType.ephemeral
    completion = requested_completion_policy or job_completion_policy
    if completion is None:
        completion = (
            CompletionPolicy.explicit
            if lifecycle == LifecycleType.daemon
            else CompletionPolicy.auto
        )
    return (lifecycle, completion)
