from __future__ import annotations

import logging
from pathlib import Path
from typing import Callable

import kuzu

from graph.connections.db_connection import DbConnection
from graph.connections.kuzu.query_result import KuzuQueryResult

logger = logging.getLogger(__name__)


class KuzuConnection(DbConnection):
    def __init__(
        self, db_path: Path | str | None, readonly: bool, create_connection: Callable[[kuzu.Database], kuzu.Connection]
    ) -> None:
        self.__db_path: Path | str | None = db_path
        self.__db: kuzu.Database | None = None
        self.__readonly = readonly
        self.__connection: kuzu.Connection | None = None
        self.__create_connection = create_connection

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

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

    def open(self) -> None:
        if not self.__connection:
            self.__db = kuzu.Database(self.__db_path, read_only=self.__readonly)
            self.__connection = self.__create_connection(self.__db)

    def close(self) -> None:
        try:
            if self.__connection:
                self.__connection.close()
        finally:
            if self.__db:
                self.__db.close()

        self._connection = None
        self._db = None

    def execute(self, query: str, parameters: dict | None = None) -> KuzuQueryResult:
        if not self.__connection:
            raise RuntimeError("Connection is not opened. Call 'open()' before executing queries.")

        return KuzuQueryResult(self.__connection.execute(query, parameters))
