from __future__ import annotations

import logging
from typing import Dict, Optional

from neo4j import Driver, GraphDatabase

from graph.connections.db_handler import DbHandler
from graph.connections.db_instance import DbInstance
from graph.connections.neo4j.db_instance import Neo4jDbInstance

logger = logging.getLogger(__name__)


class Neo4DbHandler(DbHandler):
    """
    Handles connections to a Neo4j database.
    
    Manages Neo4j instances and provides database lifecycle operations.
    Owns the lifecycle of the Neo4j driver.
    """

    def __init__(
        self,
        uri: str,
        username: str,
        password: str,
        database: str,
        timeout_seconds: Optional[int] = None,
    ):
        """
        Initialize Neo4j database handler.
        
        Args:
            uri: Neo4j connection URI (e.g., "bolt://localhost:7687").
            username: Database username.
            password: Database password.
            database: Database name.
            timeout_seconds: Query timeout in seconds.
        """
        self.uri = uri
        self.username = username
        self.password = password
        self.database = database
        self.timeout_seconds = timeout_seconds
        self.__driver: Driver | None = None

    def get_db_info(self) -> Dict:
        """Return information about the database connection."""
        return {
            "uri": self.uri,
            "username": self.username,
            "database": self.database
        }

    def _get_driver(self) -> Driver:
        """
        Get or create the Neo4j driver.
        
        Returns:
            The Neo4j driver instance.
            
        Raises:
            RuntimeError: If driver initialization fails.
        """
        if not self.__driver:
            try:
                self.__driver = GraphDatabase.driver(
                    self.uri,
                    auth=(self.username, self.password),
                )
                logger.debug(f"Neo4j driver created for {self.uri}")
            except Exception as e:
                logger.error(f"Failed to create Neo4j driver: {e}")
                raise RuntimeError(f"Failed to initialize Neo4j driver: {e}") from e
        return self.__driver

    def get_new_db_instance(self, readonly: bool) -> DbInstance:
        """
        Create a new Neo4j database instance.
        
        Args:
            readonly: Whether the instance should be read-only.
            
        Returns:
            Neo4jDbInstance ready for use.
        """
        driver = self._get_driver()
        return Neo4jDbInstance(
            driver=driver,
            database=self.database,
            readonly=readonly,
            timeout_seconds=self.timeout_seconds,
        )

    def db_exists(self) -> bool:
        """
        Check if the Neo4j database is accessible.
        
        Returns:
            True if database is reachable, False otherwise.
        """
        try:
            driver = self._get_driver()
            with driver.session(database=self.database) as session:
                session.run("RETURN 1")
            return True

        except Exception as e:
            logger.debug(f"Failed to connect to Neo4j database: {e}")
            return False

    def create_db(self) -> None:
        """Create a Neo4j database."""

        driver = self._get_driver()
        with driver.session(database="system") as session:
            try:
                session.run("CREATE DATABASE $name WAIT", {"name": self.database})
            except Exception as e:
                logger.exception("Failed to create database %s", self.database)
                raise e

        logger.info("Created temporary Neo4j DB %s", self.database)

    def delete_db(self) -> None:
        """Drop the database."""

        driver = self._get_driver()
        with driver.session(database="system") as session:
            try:
                session.run("DROP DATABASE $name IF EXISTS WAIT", {"name": self.database})
            except Exception as e:
                logger.exception("Failed to drop database %s", self.database)
                raise e

        logger.info("Dropped temporary Neo4j DB %s", self.database)

    def cleanup_handler(self) -> None:
        """
        Clean up handler resources.
        
        Closes the Neo4j driver if it was created.
        """
        if self.__driver:
            try:
                self.__driver.close()
                self.__driver = None
                logger.debug("Neo4j driver closed and cleaned up")
            except Exception as e:
                logger.exception("Failed to close Neo4j driver during cleanup")
                raise e
