from __future__ import annotations

import logging
import os
from pathlib import Path

import kuzu
from diskcache import Index

from backends.editor.utils.store import get_or_create_kuzu_db_instance_access_diskcache_dir
from graph.connections.db_instance import DbInstance
from graph.connections.internal.db_read_write_lock import DbReadWriteLock
from graph.connections.kuzu.db_connection import KuzuConnection

logger = logging.getLogger(__name__)


class KuzuDbInstance(DbInstance):
    def __init__(self, db_path: Path | str, readonly: bool) -> None:
        self.readonly = readonly

        diskcache_path = get_or_create_kuzu_db_instance_access_diskcache_dir()
        logger.debug(f"Initializing kuzu db instance access index with diskcache at {diskcache_path}.")
        self.read_write_lock = DbReadWriteLock(Index(diskcache_path))

        self.__db_path = Path(db_path)

    def open(self) -> None:
        self.read_write_lock.try_acquire_lock(self.get_db_path(), self.readonly)
        logger.debug(f"KuzuDbInstance successfully opened.")

    def close(self) -> None:
        self.read_write_lock.release_lock(self.get_db_path(), self.readonly)
        logger.debug(f"KuzuDbInstance successfully closed.")

    def get_db_path(self) -> Path:
        return self.__db_path

    def exists(self) -> bool:
        return os.path.exists(self.get_db_path())

    def get_new_conn(self, timeout_seconds: int | None = None) -> KuzuConnection:
        # We create one database instance for each query because of this issue:
        # https://github.com/kuzudb/kuzu/issues/2934
        def create_connection(db: kuzu.Database) -> kuzu.Connection:
            conn = kuzu.Connection(db)
            if timeout_seconds is not None:
                conn.execute(
                    f"""
                        CALL TIMEOUT={timeout_seconds * 1000};
                    """
                )
            return conn

        return KuzuConnection(self.get_db_path(), self.readonly, create_connection)
