"""MCP tool registrations for sub-agent lifecycle tracking."""

from mcp.server.fastmcp import Context


def register_tools(mcp, api_fn):
    @mcp.tool()
    async def spawn_subagent(
        project_id: str,
        label: str,
        parent_task_id: int | None = None,
        model: str = "",
        timeout_seconds: int = 600,
        idempotency_key: str = "",
        ctx: Context = None,
    ) -> dict | str:
        """Register a new sub-agent spawn, validating concurrency limits.

        Creates a sub-agent record and enforces project-level concurrency limits.
        If an idempotency_key is provided and a record with that key already exists,
        returns the existing record without creating a duplicate.

        Args:
            project_id: The project this sub-agent belongs to.
            label: Human-readable purpose (e.g. 'security-review', 'research-auth').
            parent_task_id: The task that spawned this sub-agent (optional).
            model: Model used for this sub-agent (optional).
            timeout_seconds: Maximum run time before kill (default 600).
            idempotency_key: Prevent duplicate spawns (optional).
        """
        return await api_fn(ctx).spawn_subagent(
            project_id, label, parent_task_id, model, timeout_seconds, idempotency_key
        )

    @mcp.tool()
    async def list_subagents(
        project_id: str | None = None,
        parent_task_id: int | None = None,
        status: str | None = None,
        ctx: Context = None,
    ) -> list[dict] | str:
        """List sub-agents with optional filters.

        Args:
            project_id: Filter by project (optional).
            parent_task_id: Filter by parent task (optional).
            status: Filter by status — 'pending', 'running', 'completed', 'failed', 'timeout', 'cancelled' (optional).
        """
        return await api_fn(ctx).list_subagents(project_id, parent_task_id, status)

    @mcp.tool()
    async def get_subagent(subagent_id: int, ctx: Context = None) -> dict | str:
        """Get a sub-agent by ID.

        Args:
            subagent_id: The numeric ID of the sub-agent.
        """
        return await api_fn(ctx).get_subagent(subagent_id)

    @mcp.tool()
    async def update_subagent(
        subagent_id: int,
        status: str | None = None,
        result_summary: str | None = None,
        tokens_used: int | None = None,
        ended_at: str | None = None,
        pid: int | None = None,
        ctx: Context = None,
    ) -> dict | str:
        """Update a sub-agent's status, result, or other fields.

        Args:
            subagent_id: The numeric ID of the sub-agent.
            status: New status — 'pending', 'running', 'completed', 'failed', 'timeout', 'cancelled' (optional).
            result_summary: Sub-agent's output or conclusion (optional).
            tokens_used: Number of tokens consumed (optional).
            ended_at: ISO timestamp when the sub-agent ended (optional).
            pid: OS process ID for timeout kill (optional).
        """
        return await api_fn(ctx).update_subagent(
            subagent_id, status, result_summary, tokens_used, ended_at, pid
        )

    @mcp.tool()
    async def kill_subagent(subagent_id: int, ctx: Context = None) -> dict | str:
        """Request timeout kill of a running sub-agent.

        Sends SIGTERM to the sub-agent process (if PID is known) and sets
        status to 'cancelled'. Only works for pending or running sub-agents.

        Args:
            subagent_id: The numeric ID of the sub-agent to kill.
        """
        return await api_fn(ctx).kill_subagent(subagent_id)

    @mcp.tool()
    async def check_subagent_limits(project_id: str, ctx: Context = None) -> dict | str:
        """Check current sub-agent usage against project limits.

        Returns current counts of running/pending sub-agents and whether
        a new sub-agent can be spawned without exceeding limits.

        Args:
            project_id: The project to check limits for.
        """
        return await api_fn(ctx).check_subagent_limits(project_id)
