"""MCP tool registrations for agent sessions."""

from mcp.server.fastmcp import Context


def register_tools(mcp, api_fn):
    @mcp.tool()
    async def create_session(
        project_id: str,
        task_id: int | None = None,
        loop_id: str = "",
        model: str = "",
        hostname: str = "",
        ctx: Context = None,
    ) -> dict | str:
        """Create a new agent session record.

        Call this when an agent loop starts a new work session. Records the
        project, associated task, loop identifier, and model in use.

        Args:
            project_id: The project identifier (org/repo format, e.g. 'pyrabbit/peon-mcp').
            task_id: The task this session is working on (optional).
            loop_id: The loop instance identifier (optional).
            model: The model name in use (optional).
            hostname: The hostname of the machine running the session (optional).
        """
        return await api_fn(ctx).create_session(project_id, task_id, loop_id, model, hostname)

    @mcp.tool()
    async def update_session(
        session_id: int,
        status: str | None = None,
        exit_code: int | None = None,
        tokens_input: int | None = None,
        tokens_output: int | None = None,
        cost_usd: float | None = None,
        error_summary: str | None = None,
        ended_at: str | None = None,
        ctx: Context = None,
    ) -> dict | str:
        """Update an agent session (e.g. to end it, record token usage).

        Args:
            session_id: The numeric session ID.
            status: New status — 'running', 'completed', 'failed', or 'timeout'.
            exit_code: Process exit code (optional).
            tokens_input: Total input tokens consumed (optional).
            tokens_output: Total output tokens produced (optional).
            cost_usd: Total API cost in USD (optional).
            error_summary: Brief description of the error if failed (optional).
            ended_at: ISO timestamp when the session ended (optional).
        """
        return await api_fn(ctx).update_session(
            session_id, status, exit_code, tokens_input, tokens_output,
            cost_usd, error_summary, ended_at,
        )

    @mcp.tool()
    async def add_session_event(
        session_id: int,
        event_type: str,
        tool_name: str = "",
        detail: str = "",
        tool_args: dict | None = None,
        ctx: Context = None,
    ) -> dict | str:
        """Append an event to an agent session.

        Use this to record significant tool calls, errors, or checkpoints
        during a session for later analysis.

        Args:
            session_id: The numeric session ID.
            event_type: One of 'tool_call', 'tool_result', 'file_read', 'file_write', 'git_op', 'error', 'info', 'message', 'checkpoint', 'start'.
            tool_name: Name of the tool called (if event_type is tool_call/tool_result).
            detail: Additional detail or payload for the event.
            tool_args: Structured tool call arguments as a JSON object (optional).
        """
        return await api_fn(ctx).add_session_event(session_id, event_type, tool_name, detail, tool_args)

    @mcp.tool()
    async def list_sessions(
        project_id: str | None = None,
        task_id: int | None = None,
        status: str | None = None,
        ctx: Context = None,
    ) -> list[dict] | str:
        """List agent sessions with optional filters.

        Args:
            project_id: Filter to sessions for a specific project (optional).
            task_id: Filter to sessions for a specific task (optional).
            status: Filter by status — 'running', 'completed', 'failed', 'timeout' (optional).
        """
        return await api_fn(ctx).list_sessions(project_id, task_id, status)

    @mcp.tool()
    async def get_session(session_id: int, ctx: Context = None) -> dict | str:
        """Get a single agent session with all its events.

        Args:
            session_id: The numeric ID of the session.
        """
        return await api_fn(ctx).get_session(session_id)

    @mcp.tool()
    async def session_stats(project_id: str, ctx: Context = None) -> dict | str:
        """Get aggregate statistics for a project's agent sessions.

        Returns total counts, success rate, average duration, token totals,
        total cost, and sessions broken down by model.

        Args:
            project_id: The project identifier (org/repo format, e.g. 'pyrabbit/peon-mcp').
        """
        return await api_fn(ctx).session_stats(project_id)
