from __future__ import annotations

import json
import logging
import os
import tempfile
from pathlib import Path
from typing import Dict

import kuzu
from dataiku import Folder

from graph.connections.db_handler import DbHandler
from graph.connections.db_instance import DbInstance
from graph.connections.kuzu.db_instance import KuzuDbInstance

logger = logging.getLogger(__name__)


class KuzuLocalDbHandler(DbHandler):
    def __init__(self, db_path: Path | str, should_db_persist: bool = True):
        self.db_path = db_path
        self.should_db_persist = should_db_persist

        root_path = Path(self.db_path).parent
        if not root_path.exists():
            root_path.mkdir(parents=True, exist_ok=True)

    def get_db_info(self) -> Dict:
        return {"db_path": str(self.db_path)}

    def get_new_db_instance(self, readonly: bool) -> DbInstance:
        return KuzuDbInstance(db_path=self.db_path, readonly=readonly)

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

    def create_db(self) -> None:
        if self.db_exists():
            raise RuntimeError("Cannot create a kuzu database, it already exists at this path.")

        conn = kuzu.Database(self.db_path)
        conn.close()

    def delete_db(self) -> None:
        os.unlink(self.db_path)

    def cleanup_handler(self) -> None:
        if not self.should_db_persist and self.db_exists():
            self.delete_db()


class KuzuLocalReplicaDbHandler(DbHandler):
    def __init__(self, db_path: Path | str, db_folder: Folder):
        self.db_path = Path(db_path)
        self.db_folder = db_folder
        self.temp_root = Path(tempfile.gettempdir())

        self.subdir = self.db_path.parent
        self.__temp_dir = (self.temp_root / self.subdir).resolve()
        self.__temp_dir.mkdir(parents=True, exist_ok=True)
        self.local_db_file_path = self.__temp_dir / self.db_path.name

    def get_db_info(self) -> Dict:
        return {"db_path": str(self.db_path), "db_folder": self.db_folder}

    def get_new_db_instance(self, readonly: bool) -> DbInstance:
        build_info_name = "buildInfo.json"
        remote_build_info_path = self.subdir / build_info_name
        local_build_info_path = self.__temp_dir / build_info_name

        need_download = self._sync_build_info(local_build_info_path, remote_build_info_path, self.db_folder)

        if need_download or not self.local_db_file_path.exists():
            logger.info(f"Downloading {str(self.db_path)} to {str(self.local_db_file_path)}...")
            with (
                self.db_folder.get_download_stream(str(self.db_path)) as remote_file,
                self.local_db_file_path.open("wb") as out,
            ):
                out.write(remote_file.read())
        else:
            logger.info(f"Using cached DB at {str(self.local_db_file_path)} (up to date).")

        return KuzuDbInstance(db_path=self.local_db_file_path, readonly=readonly)

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

    def create_db(self) -> None:
        raise RuntimeError("Cannot create a kuzu replica database.")

    def delete_db(self) -> None:
        raise RuntimeError("Cannot delete a kuzu replica database.")

    def cleanup_handler(self) -> None:
        os.unlink(self.local_db_file_path)

    @staticmethod
    def _safe_load_json(path: Path) -> dict | None:
        """Load JSON from `path`; return None if the file doesn't exist or is invalid."""
        if not path.exists():
            return None
        try:
            with path.open("r", encoding="utf-8") as f:
                return json.load(f)  # type: ignore
        except (OSError, json.JSONDecodeError):
            return None

    @staticmethod
    def _sync_build_info(local_path: Path, remote_path: Path, db_folder: Folder) -> bool:
        """
        Ensure local `buildInfo.json` matches remote.
        Returns True if the DB should be (re)downloaded (missing or changed).
        """
        with db_folder.get_download_stream(str(remote_path)) as remote_file:
            remote_bytes = remote_file.read()
        try:
            remote_info = json.loads(remote_bytes)
            logger.debug(f"Folder epoc {remote_info}")
        except json.JSONDecodeError:
            remote_info = None

        local_info = KuzuLocalReplicaDbHandler._safe_load_json(local_path)
        logger.debug(f"Local epoc {local_info} {str(local_path)}")

        changed = (local_info is None) or (remote_info is None) or (remote_info != local_info)

        try:
            with local_path.open("wb") as f:
                f.write(remote_bytes)
        except OSError:
            changed = True

        return changed
