from __future__ import annotations

import logging

from neo4j import Driver, 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 = self.__session.run(query, parameters or {}, timeout=self.__timeout_seconds) # type: ignore
            records = list(result)  # Consume all records
            return Neo4jQueryResult(records)
        except Exception:
            logger.exception(f"Failed to execute query: {query}")
            raise
