"""Knowledge graph 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 KnowledgeResource:
    """Sync knowledge graph operations."""

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

    def ingest(
        self,
        content: str,
        *,
        entity_id: str | None = None,
        metadata: dict[str, Any] | None = None,
        source: str | None = None,
    ) -> APIResponse[Any]:
        """Add content to the knowledge graph."""
        return self._client.request(
            "POST",
            "/api/v1/knowledge/ingest",
            json={"content": content, "entity_id": entity_id, "metadata": metadata, "source": source},
        )

    def query(
        self,
        query: str,
        *,
        entity_id: str | None = None,
        depth: int = 2,
        limit: int = 20,
    ) -> APIResponse[Any]:
        """Query the knowledge graph."""
        return self._client.request(
            "POST",
            "/api/v1/knowledge/query",
            json={"query": query, "entity_id": entity_id, "depth": depth, "limit": limit},
        )

    def entities(self, *, entity_type: str | None = None, limit: int = 50) -> APIResponse[Any]:
        """List entities in the knowledge graph."""
        params: dict[str, Any] = {"limit": limit}
        if entity_type:
            params["entity_type"] = entity_type
        return self._client.request("GET", "/api/v1/knowledge/entities", params=params)

    def visualize(self, query: str, *, depth: int = 2) -> APIResponse[Any]:
        """Get visualization data for the knowledge graph."""
        return self._client.request(
            "POST",
            "/api/v1/knowledge/visualize",
            json={"query": query, "depth": depth},
        )


class AsyncKnowledgeResource:
    """Async knowledge graph operations."""

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

    async def ingest(
        self,
        content: str,
        *,
        entity_id: str | None = None,
        metadata: dict[str, Any] | None = None,
        source: str | None = None,
    ) -> APIResponse[Any]:
        """Add content to the knowledge graph."""
        return await self._client.request(
            "POST",
            "/api/v1/knowledge/ingest",
            json={"content": content, "entity_id": entity_id, "metadata": metadata, "source": source},
        )

    async def query(
        self,
        query: str,
        *,
        entity_id: str | None = None,
        depth: int = 2,
        limit: int = 20,
    ) -> APIResponse[Any]:
        """Query the knowledge graph."""
        return await self._client.request(
            "POST",
            "/api/v1/knowledge/query",
            json={"query": query, "entity_id": entity_id, "depth": depth, "limit": limit},
        )

    async def entities(self, *, entity_type: str | None = None, limit: int = 50) -> APIResponse[Any]:
        """List entities in the knowledge graph."""
        params: dict[str, Any] = {"limit": limit}
        if entity_type:
            params["entity_type"] = entity_type
        return await self._client.request("GET", "/api/v1/knowledge/entities", params=params)

    async def visualize(self, query: str, *, depth: int = 2) -> APIResponse[Any]:
        """Get visualization data for the knowledge graph."""
        return await self._client.request(
            "POST",
            "/api/v1/knowledge/visualize",
            json={"query": query, "depth": depth},
        )
