"""Graph traversal 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 GraphResource:
    """Sync graph traversal operations."""

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

    def traverse(
        self,
        *,
        entity_name: str | None = None,
        entity_type: str | None = None,
        relationship_types: list[str] | None = None,
        depth: int = 2,
    ) -> APIResponse[Any]:
        """Traverse the knowledge graph from a starting entity."""
        body: dict[str, Any] = {"depth": depth}
        if entity_name:
            body["entity_name"] = entity_name
        if entity_type:
            body["entity_type"] = entity_type
        if relationship_types:
            body["relationship_types"] = relationship_types
        return self._client.request("POST", "/api/v1/graph/traverse", json=body)

    def temporal(
        self,
        entity_id: str,
        as_of_date: str,
    ) -> APIResponse[Any]:
        """Get temporal view of entity relationships at a specific date."""
        return self._client.request(
            "POST",
            "/api/v1/graph/temporal",
            json={"entity_id": entity_id, "as_of_date": as_of_date},
        )

    def neighbors(self, entity_id: str) -> APIResponse[Any]:
        """Get immediate neighbors of an entity."""
        return self._client.request("GET", f"/api/v1/graph/neighbors/{entity_id}")


class AsyncGraphResource:
    """Async graph traversal operations."""

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

    async def traverse(
        self,
        *,
        entity_name: str | None = None,
        entity_type: str | None = None,
        relationship_types: list[str] | None = None,
        depth: int = 2,
    ) -> APIResponse[Any]:
        """Traverse the knowledge graph from a starting entity."""
        body: dict[str, Any] = {"depth": depth}
        if entity_name:
            body["entity_name"] = entity_name
        if entity_type:
            body["entity_type"] = entity_type
        if relationship_types:
            body["relationship_types"] = relationship_types
        return await self._client.request("POST", "/api/v1/graph/traverse", json=body)

    async def temporal(
        self,
        entity_id: str,
        as_of_date: str,
    ) -> APIResponse[Any]:
        """Get temporal view of entity relationships at a specific date."""
        return await self._client.request(
            "POST",
            "/api/v1/graph/temporal",
            json={"entity_id": entity_id, "as_of_date": as_of_date},
        )

    async def neighbors(self, entity_id: str) -> APIResponse[Any]:
        """Get immediate neighbors of an entity."""
        return await self._client.request("GET", f"/api/v1/graph/neighbors/{entity_id}")
