"""Unified async graph client — abstracts Neo4j, Memgraph, and FalkorDB."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from koan.graph.config import GraphConfig


@dataclass
class GraphRecord:
    """A single result row from a Cypher query."""

    data: dict[str, Any] = field(default_factory=dict)


class GraphClient(ABC):
    """Unified async interface over property graph databases.

    All three supported databases (Neo4j, Memgraph, FalkorDB) speak Cypher.
    This ABC normalizes the transport difference (Bolt vs Redis protocol).

    Usage::

        async with create_graph_client(config) as client:
            records = await client.execute("MATCH (n) RETURN n LIMIT 10")
            await client.execute_write("CREATE (n:Test {id: $id})", {"id": "1"})
    """

    @abstractmethod
    async def connect(self) -> None:
        """Open the database connection."""
        ...

    @abstractmethod
    async def close(self) -> None:
        """Close the database connection."""
        ...

    @abstractmethod
    async def execute(
        self,
        cypher: str,
        params: dict[str, Any] | None = None,
    ) -> list[GraphRecord]:
        """Execute a read query and return results."""
        ...

    @abstractmethod
    async def execute_write(
        self,
        cypher: str,
        params: dict[str, Any] | None = None,
    ) -> list[GraphRecord]:
        """Execute a write query and return results."""
        ...

    @abstractmethod
    async def health_check(self) -> bool:
        """Return True if the database is reachable."""
        ...

    async def __aenter__(self) -> GraphClient:
        await self.connect()
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.close()


class Neo4jClient(GraphClient):
    """Async client for Neo4j and Memgraph (both use Bolt protocol)."""

    def __init__(self, config: GraphConfig) -> None:
        self.config = config
        self._driver: Any = None

    async def connect(self) -> None:
        from neo4j import AsyncGraphDatabase  # type: ignore[import-untyped]

        auth = (self.config.username, self.config.password) if self.config.username else None
        self._driver = AsyncGraphDatabase.driver(self.config.uri, auth=auth)

    async def close(self) -> None:
        if self._driver:
            await self._driver.close()

    async def execute(
        self,
        cypher: str,
        params: dict[str, Any] | None = None,
    ) -> list[GraphRecord]:
        async with self._driver.session(database=self.config.database) as session:
            result = await session.run(cypher, params or {})
            records = [GraphRecord(data=dict(r)) async for r in result]
            return records

    async def execute_write(
        self,
        cypher: str,
        params: dict[str, Any] | None = None,
    ) -> list[GraphRecord]:
        async with self._driver.session(database=self.config.database) as session:

            async def _tx_work(tx: Any) -> list[dict[str, Any]]:
                result = await tx.run(cypher, params or {})
                return [dict(r) async for r in result]

            rows = await session.execute_write(_tx_work)
            return [GraphRecord(data=row) for row in rows]

    async def health_check(self) -> bool:
        try:
            await self._driver.verify_connectivity()
        except Exception:  # noqa: BLE001
            return False
        return True


class FalkorDBClient(GraphClient):
    """Async client for FalkorDB (Redis protocol, Cypher queries)."""

    def __init__(self, config: GraphConfig) -> None:
        self.config = config
        self._db: Any = None
        self._graph: Any = None

    async def connect(self) -> None:
        from falkordb.asyncio import FalkorDB  # type: ignore[import-untyped]

        self._db = FalkorDB(host=self._parse_host(), port=self._parse_port())
        self._graph = self._db.select_graph(self.config.database)

    async def close(self) -> None:
        if self._db:
            await self._db.aclose()

    async def execute(
        self,
        cypher: str,
        params: dict[str, Any] | None = None,
    ) -> list[GraphRecord]:
        result = await self._graph.query(cypher, params=params or {})
        # FalkorDB headers are [column_type, column_name] pairs
        headers = [h[1] if isinstance(h, list) else h for h in result.header]
        return [GraphRecord(data=dict(zip(headers, row, strict=False))) for row in result.result_set]

    async def execute_write(
        self,
        cypher: str,
        params: dict[str, Any] | None = None,
    ) -> list[GraphRecord]:
        return await self.execute(cypher, params)

    async def health_check(self) -> bool:
        try:
            await self._graph.query("RETURN 1")
        except Exception:  # noqa: BLE001
            return False
        return True

    def _parse_host(self) -> str:
        return self.config.uri.replace("redis://", "").split(":")[0]

    def _parse_port(self) -> int:
        parts = self.config.uri.replace("redis://", "").split(":")
        return int(parts[1]) if len(parts) > 1 else 6379


def create_graph_client(config: GraphConfig) -> GraphClient:
    """Factory — returns the appropriate client for the configured provider."""
    if config.provider == "falkordb":
        return FalkorDBClient(config)
    # Neo4j and Memgraph both use Bolt protocol
    return Neo4jClient(config)
