"""API client methods for agent sessions."""

from __future__ import annotations

from typing import Any

from peon_mcp.common.base_client import BaseAPIClient


class SessionClient(BaseAPIClient):
    """Agent session API client methods."""

    async def create_session(
        self,
        project_id: str,
        task_id: int | None = None,
        loop_id: str = "",
        model: str = "",
        hostname: str = "",
    ) -> dict | str:
        return await self._request(
            "POST",
            "/api/sessions",
            json={
                "project_id": project_id,
                "task_id": task_id,
                "loop_id": loop_id,
                "model": model,
                "hostname": hostname,
            },
        )

    async def update_session(
        self,
        session_id: int,
        task_id: int | None = None,
        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,
    ) -> dict | str:
        payload: dict[str, Any] = {}
        if task_id is not None:
            payload["task_id"] = task_id
        if status is not None:
            payload["status"] = status
        if exit_code is not None:
            payload["exit_code"] = exit_code
        if tokens_input is not None:
            payload["tokens_input"] = tokens_input
        if tokens_output is not None:
            payload["tokens_output"] = tokens_output
        if cost_usd is not None:
            payload["cost_usd"] = cost_usd
        if error_summary is not None:
            payload["error_summary"] = error_summary
        if ended_at is not None:
            payload["ended_at"] = ended_at
        return await self._request("PUT", f"/api/sessions/{session_id}", json=payload)

    async def add_session_event(
        self,
        session_id: int,
        event_type: str,
        tool_name: str = "",
        detail: str = "",
        tool_args: dict | None = None,
    ) -> dict | str:
        payload: dict[str, Any] = {
            "session_id": session_id,
            "event_type": event_type,
            "tool_name": tool_name,
            "detail": detail,
        }
        if tool_args is not None:
            payload["tool_args"] = tool_args
        return await self._request(
            "POST",
            f"/api/sessions/{session_id}/events",
            json=payload,
        )

    async def list_sessions(
        self,
        project_id: str | None = None,
        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 task_id is not None:
            params["task_id"] = task_id
        if status is not None:
            params["status"] = status
        return await self._paginate_all("/api/sessions", params=params)

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

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