"""Agent resource for the Cerebe SDK."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from .._response import APIResponse

if TYPE_CHECKING:
    from .._client import AsyncCerebe, Cerebe


class AgentResource:
    """Sync agent operations — trace ingestion and working memory."""

    def __init__(self, client: Cerebe) -> None:
        self._client = client

    def ingest_trace(
        self,
        content: str,
        session_id: str,
        *,
        metadata: dict[str, Any] | None = None,
    ) -> APIResponse[Any]:
        """Ingest an agent execution trace."""
        body: dict[str, Any] = {
            "content": content,
            "session_id": session_id,
            "memory_type": "execution_history",
        }
        if metadata:
            body["metadata"] = metadata
        return self._client.request("POST", "/api/v1/memory/store", json=body)

    def set_working_memory(
        self,
        content: str,
        session_id: str,
        *,
        ttl_seconds: int | None = None,
    ) -> APIResponse[Any]:
        """Set working memory for the current session."""
        body: dict[str, Any] = {
            "content": content,
            "session_id": session_id,
            "memory_type": "working",
        }
        if ttl_seconds is not None:
            body["metadata"] = {"ttl_seconds": ttl_seconds}
        return self._client.request("POST", "/api/v1/memory/store", json=body)

    def get_working_memory(self, session_id: str) -> APIResponse[Any]:
        """Get working memory for a session."""
        return self._client.request(
            "GET",
            f"/api/v1/memory/session/{session_id}",
            params={"memory_types": "working", "limit": "10"},
        )


class AsyncAgentResource:
    """Async agent operations — trace ingestion and working memory."""

    def __init__(self, client: AsyncCerebe) -> None:
        self._client = client

    async def ingest_trace(
        self,
        content: str,
        session_id: str,
        *,
        metadata: dict[str, Any] | None = None,
    ) -> APIResponse[Any]:
        """Ingest an agent execution trace."""
        body: dict[str, Any] = {
            "content": content,
            "session_id": session_id,
            "memory_type": "execution_history",
        }
        if metadata:
            body["metadata"] = metadata
        return await self._client.request("POST", "/api/v1/memory/store", json=body)

    async def set_working_memory(
        self,
        content: str,
        session_id: str,
        *,
        ttl_seconds: int | None = None,
    ) -> APIResponse[Any]:
        """Set working memory for the current session."""
        body: dict[str, Any] = {
            "content": content,
            "session_id": session_id,
            "memory_type": "working",
        }
        if ttl_seconds is not None:
            body["metadata"] = {"ttl_seconds": ttl_seconds}
        return await self._client.request("POST", "/api/v1/memory/store", json=body)

    async def get_working_memory(self, session_id: str) -> APIResponse[Any]:
        """Get working memory for a session."""
        return await self._client.request(
            "GET",
            f"/api/v1/memory/session/{session_id}",
            params={"memory_types": "working", "limit": "10"},
        )
