"""Agent lifecycle setup and teardown helpers.

Extracted from agent.py to reduce file size.
Spec source: vectl step ``shrink-agent-py`` in ``gate-review-debt`` phase.
"""

# @invar:allow shell_result: Lifecycle helpers return Result types
# @invar:allow forbidden_import: Shell module needs os for PID/env, pathlib for paths
# @shell_orchestration: Lifecycle setup delegates to state I/O and signal registration

from __future__ import annotations

import os
import signal
from typing import TYPE_CHECKING, Any

from returns.result import Failure, Result, Success

from anima.core.lifecycle import LifecycleEvent
from anima.core.schemas import (
    AgentState,
    AgentStateData,
    AnimaError,
    AnimaTimeoutError,
    InternalError,
)
from anima.shell.state import (
    append_log,
    get_state_dir,
    read_state,
    transition_state,
    write_agent_pid,
)
from anima.shell.timeout_watchdog import (
    AsyncTimeoutWatchdog,
    create_async_timeout_watchdog,
)

if TYPE_CHECKING:
    from anima.core.schemas import CapabilityToken, JobSpec
    from anima.shell.mcp_agent import AgentMCPServer


# @shell_orchestration: Registers OS signal handlers (I/O boundary)
def setup_signal_handlers(
    interrupted_setter: Any,
    mcp_server_getter: Any,
) -> None:
    """Set up signal handlers for graceful shutdown.

    Args:
        interrupted_setter: Callable to set interrupted flag.
        mcp_server_getter: Callable returning MCP server or None.
    """

    def handle_signal(signum: int, frame: object) -> None:
        interrupted_setter()
        server = mcp_server_getter()
        if server is not None:
            event = server.get_interrupt_event()
            event.set()

    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)


# @shell_orchestration: State transition and log I/O on timeout
def on_timeout(
    instance_id: str,
    timeout_fired_setter: Any,
    interrupted_setter: Any,
    mcp_server: Any,
) -> None:
    """Handle timeout by transitioning agent to failed state.

    Per docs/INTERFACES.md Section H.4:
    - Timeout is a hard wall-clock limit
    - Transition agent to 'failed' state immediately
    - No grace finalization for timeout

    Args:
        instance_id: The agent that timed out.
        timeout_fired_setter: Callable to set timeout_fired flag.
        interrupted_setter: Callable to set interrupted flag.
        mcp_server: MCP server instance or None.
    """
    timeout_fired_setter()
    transition_state(
        instance_id,
        LifecycleEvent.timeout,
        error="Agent timed out (wall-clock timeout exceeded)",
    )
    append_log(
        instance_id,
        "timeout",
        {"message": "Agent timed out after configured duration"},
    )
    interrupted_setter()
    if mcp_server is not None:
        try:
            event = mcp_server.get_interrupt_event()
            event.set()
        except RuntimeError:
            pass


def configure_run_environment(
    instance_id: str,
    state: AgentStateData,
    capability_token: "CapabilityToken | None",
) -> None:
    """Configure process environment required by run initialization.

    Args:
        instance_id: The agent's unique identifier.
        state: Agent state data with trace_id.
        capability_token: Optional capability token.
    """
    os.environ["ANIMA_AGENT_ID"] = instance_id
    if "ANIMA_STATE" not in os.environ:
        os.environ["ANIMA_STATE"] = str(get_state_dir())
    if state.trace_id:
        os.environ["ANIMA_TRACE_ID"] = state.trace_id
    if capability_token is not None:
        os.environ["ANIMA_CAPABILITY_TOKEN_JSON"] = capability_token.model_dump_json()


# @shell_complexity: State initialization with PID and socket setup
async def initialize_run_state(
    instance_id: str,
    capability_token: "CapabilityToken | None",
    setup_signals_fn: Any,
    setup_watchdog_fn: Any,
) -> Result[AgentStateData, AnimaError]:
    """Initialize persisted run state, PID, and watchdog.

    Spec source: medium-fixes.reduce-agent-size target (b).

    Args:
        instance_id: The agent's unique identifier.
        capability_token: Optional capability token.
        setup_signals_fn: Callable to set up signal handlers.
        setup_watchdog_fn: Async callable to set up timeout watchdog.

    Returns:
        Success with agent state data, or Failure with error.
    """
    state_result = read_state(instance_id)
    if isinstance(state_result, Failure):
        return Failure(
            InternalError(f"Failed to read agent state: {state_result.failure()}")
        )
    state = state_result.unwrap()

    configure_run_environment(instance_id, state, capability_token)

    pid_result = write_agent_pid(instance_id, os.getpid())
    if isinstance(pid_result, Failure):
        return Failure(InternalError(f"Failed to write PID: {pid_result.failure()}"))
    if state.state != AgentState.running:
        trans_result = transition_state(instance_id, LifecycleEvent.start)
        if isinstance(trans_result, Failure):
            return Failure(
                InternalError(
                    f"Failed to transition to running: {trans_result.failure()}"
                )
            )

    setup_signals_fn()
    timeout_result = await setup_watchdog_fn()
    if isinstance(timeout_result, Failure):
        return timeout_result

    return Success(state)


# @shell_complexity: Timeout parsing with optional watchdog creation
async def setup_timeout_watchdog(
    timeout: str | None,
    job: "JobSpec | None",
    instance_id: str,
    on_timeout_fn: Any,
) -> tuple[Result[None, AnimaError], AsyncTimeoutWatchdog | None]:
    """Initialize and start timeout watchdog when timeout is configured.

    Args:
        timeout: Optional timeout string.
        job: Optional job spec with timeout.
        instance_id: Agent's unique identifier.
        on_timeout_fn: Callback for timeout events.

    Returns:
        Tuple of (result, watchdog_instance_or_none).
    """
    timeout_str = timeout
    if timeout_str is None and job is not None:
        timeout_str = job.timeout
    if timeout_str is None:
        return Success(None), None

    watchdog_result = await create_async_timeout_watchdog(
        timeout_str,
        instance_id,
        on_timeout_fn,
    )
    if isinstance(watchdog_result, Failure):
        transition_state(
            instance_id,
            LifecycleEvent.fail,
            error=f"Invalid timeout format: {timeout_str}",
        )
        return Failure(
            AnimaTimeoutError(f"Invalid timeout format: {timeout_str}")
        ), None

    watchdog = watchdog_result.unwrap()
    start_result = await watchdog.start()
    if isinstance(start_result, Failure):
        transition_state(
            instance_id,
            LifecycleEvent.fail,
            error=str(start_result.failure()),
        )
        return Failure(start_result.failure()), None
    return Success(None), watchdog


# @invar:allow shell_result: Pass-through returns caller's Result
# @shell_orchestration: Cancels watchdog and stops MCP server (I/O cleanup)
async def finalize_run(
    result: Any,
    timeout_watchdog: AsyncTimeoutWatchdog | None,
    mcp_server: "AgentMCPServer | None",
) -> Any:
    """Finalize run cleanup and return the original result.

    Spec source: medium-fixes.reduce-agent-size target (c).

    Args:
        result: The run result to pass through.
        timeout_watchdog: Optional watchdog to cancel.
        mcp_server: Optional MCP server to stop.

    Returns:
        The original result, unchanged.
    """
    if timeout_watchdog is not None:
        try:
            await timeout_watchdog.cancel()
        except RuntimeError:
            pass
    if mcp_server is not None:
        await mcp_server.stop()
    return result


__all__ = [
    "configure_run_environment",
    "finalize_run",
    "initialize_run_state",
    "on_timeout",
    "setup_signal_handlers",
    "setup_timeout_watchdog",
]
