from __future__ import annotations

from pathlib import Path
from typing import Callable

import kuzu


class KuzuConnectionContextManager:
    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) -> KuzuConnectionContextManager:
        self.open()
        return self

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

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

        return self.__connection

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

        self._connection = None
        self._db = None

    @property
    def connection(self) -> kuzu.Connection:
        if not self.__connection:
            raise RuntimeError(
                "KuzuConnectionContextManager needs to be opened before accessing underlying connection."
            )

        return self.__connection
