from __future__ import annotations

import logging

from neo4j import Driver, Result, Session

from graph.connections.db_connection import DbConnection
from graph.connections.neo4j.query_result import Neo4jQueryResult

logger = logging.getLogger(__name__)


class Neo4jConnection(DbConnection):
    def __init__(
        self,
        driver: Driver,
        database: str,
        readonly: bool,
        timeout_seconds: int | None = None,
    ) -> None:
        """
        Initialize Neo4j connection wrapper.

        Args:
            driver: Neo4j Driver instance.
            database: Name of the database.
            readonly: Whether connection should be read-only.
            timeout_seconds: Query timeout in seconds.
        """
        self.__driver = driver
        self.__database = database
        self.__readonly = readonly
        self.__timeout_seconds = timeout_seconds
        self.__session: Session | None = None

    def __enter__(self) -> Neo4jConnection:
        self.open()
        return self

    def __exit__(self, exception_type, exception_value, exception_traceback) -> None:
        self.close()

    def open(self) -> None:
        """Open a Neo4j session."""
        if not self.__session:
            self.__session = self.__driver.session(
                default_access_mode="READ" if self.__readonly else "WRITE", database=self.__database
            )

    def close(self) -> None:
        """Close the Neo4j session."""
        try:
            if self.__session:
                self.__session.close()
        except Exception:
            logger.exception("Failed to close Neo4j session")
        finally:
            self.__session = None

    def execute(self, query: str, parameters: dict | None = None) -> Neo4jQueryResult:
        """
        Execute a Cypher query against the Neo4j database.

        Args:
            query: Cypher query string.
            parameters: Query parameters dictionary.

        Returns:
            Neo4jQueryResult wrapping the query results.

        Raises:
            RuntimeError: If connection is not opened.
        """
        if not self.__session:
            raise RuntimeError("Connection is not opened. Call 'open()' before executing queries.")

        try:
            result: Result = self.__session.run(query, parameters or {}, timeout=self.__timeout_seconds)  # type: ignore
            return Neo4jQueryResult(result)
        except Exception:
            logger.exception(f"Failed to execute query: {query}")
            raise
