"""API client methods for sub-agents."""

from __future__ import annotations

from typing import Any

from peon_mcp.common.base_client import BaseAPIClient


class SubAgentClient(BaseAPIClient):
    """Sub-agent API client methods."""

    async def spawn_subagent(
        self,
        project_id: str,
        label: str,
        parent_task_id: int | None = None,
        model: str = "",
        timeout_seconds: int = 600,
        idempotency_key: str = "",
    ) -> dict | str:
        body: dict[str, Any] = {
            "project_id": project_id,
            "label": label,
            "model": model,
            "timeout_seconds": timeout_seconds,
            "idempotency_key": idempotency_key,
        }
        if parent_task_id is not None:
            body["parent_task_id"] = parent_task_id
        return await self._request("POST", "/api/subagents", json=body)

    async def list_subagents(
        self,
        project_id: str | None = None,
        parent_task_id: int | None = None,
        status: str | None = None,
    ) -> list[dict] | str:
        params: dict[str, Any] = {}
        if project_id is not None:
            params["project_id"] = project_id
        if parent_task_id is not None:
            params["parent_task_id"] = parent_task_id
        if status is not None:
            params["status"] = status
        return await self._paginate_all("/api/subagents", params=params)

    async def get_subagent(self, subagent_id: int) -> dict | str:
        return await self._request("GET", f"/api/subagents/{subagent_id}")

    async def update_subagent(
        self,
        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,
    ) -> dict | str:
        body: dict[str, Any] = {}
        if status is not None:
            body["status"] = status
        if result_summary is not None:
            body["result_summary"] = result_summary
        if tokens_used is not None:
            body["tokens_used"] = tokens_used
        if ended_at is not None:
            body["ended_at"] = ended_at
        if pid is not None:
            body["pid"] = pid
        return await self._request("PUT", f"/api/subagents/{subagent_id}", json=body)

    async def delete_subagent(self, subagent_id: int) -> dict | str:
        return await self._request("DELETE", f"/api/subagents/{subagent_id}")

    async def kill_subagent(self, subagent_id: int) -> dict | str:
        return await self._request("POST", f"/api/subagents/{subagent_id}/kill")

    async def check_subagent_limits(self, project_id: str) -> dict | str:
        return await self._request(
            "GET", "/api/subagents/limits", params={"project_id": project_id}
        )
