from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Optional, cast

from dataiku import Folder
from dataiku.core.base import is_container_exec
from typing_extensions import TypedDict

from graph.connections.neo4j.preset import Neo4jPresetConfig
from recipes.build_graph.assets import (
    BUILD_INFO_JSON_FILENAME,
    BUILT_GRAPHS_PREFIX,
    CONFIGURATION_JSON_FILENAME,
    DATABASE_JSON_FILENAME,
    DB_KUZU_FILENAME,
)
from solutions.graph.models import DbType, GraphMetadataSnapshot

logger = logging.getLogger(__name__)


class GraphBuildInfo(TypedDict):
    epoch_ms: int


class AbstractConnectionInfo(ABC):
    """Abstract base class for connection information.

    Defines the interface for retrieving connection details for different
    database backends.
    """

    def __init__(self) -> None:
        """Initializes the base connection info."""
        super().__init__()

    @property
    @abstractmethod
    def type(self) -> DbType:
        """The type of the database (e.g., 'kuzu', 'neo4j')."""
        pass

    @property
    @abstractmethod
    def uid(self) -> str:
        """The unique identifier for this connection (snapshot ID or DB name)."""
        pass

    @property
    @abstractmethod
    def db_folder(self) -> Folder | None:
        """The DSS folder where the database or its metadata is stored."""
        pass

    @property
    def configuration_path(self) -> Path:
        """The relative path to the configuration JSON file within the folder."""
        return Path(BUILT_GRAPHS_PREFIX) / self.uid / CONFIGURATION_JSON_FILENAME

    @property
    def build_info_path(self) -> Path:
        """The relative path to the build info JSON file within the folder."""
        return Path(BUILT_GRAPHS_PREFIX) / self.uid / BUILD_INFO_JSON_FILENAME

    def get_build_info(self) -> GraphBuildInfo | None:
        """Read the build information from the DSS folder.

        Returns:
            Optional[GraphBuildInfo]: The validated build info, or None if no folder exists.
        """
        if self.db_folder:
            build_info = self.db_folder.read_json(self.build_info_path)
            if "epoch_ms" not in build_info:
                raise ValueError(f"Missing 'epoch_ms' in build info for snapshot {self.uid}")
            return cast(GraphBuildInfo, build_info)
        return None

    def get_configuration(self) -> GraphMetadataSnapshot | None:
        """Read the graph configuration from the DSS folder.

        Returns:
            Optional[GraphMetadataSnapshot]: The configuration, or None if no folder exists.
        """
        if self.db_folder:
            configuration = self.db_folder.read_json(self.configuration_path)
            return cast(GraphMetadataSnapshot, configuration)
        return None

class KuzuConnectionInfo(AbstractConnectionInfo):
    """Connection information for a Kuzu (local) database.

    Attributes:
        db_folder (Folder): The Dataiku folder containing the Kuzu database.
        snapshot_id (str): The unique identifier for the graph snapshot.
    """

    def __init__(self, db_folder: Folder, snapshot_id: str) -> None:
        """Initializes Kuzu connection info.

        Args:
            db_folder (Folder): The DSS folder containing the database.
            snapshot_id (str): Unique identifier for the snapshot.
        """
        self._db_folder = db_folder
        self.snapshot_id = snapshot_id

    @property
    def type(self) -> DbType:
        """The type of the database ('kuzu')."""
        return "kuzu"

    @property
    def uid(self) -> str:
        """The unique identifier for this connection (snapshot ID)."""
        return self.snapshot_id
    
    @property
    def db_folder(self) -> Folder:
        """The DSS folder where the Kuzu database is stored."""
        return self._db_folder
    
    @property
    def is_remote(self) -> bool:
        """Check if the database is stored on a remote filesystem (S3, GCS, etc.).

        Returns:
            bool: True if the folder is not local filesystem or in container exec.
        """
        db_folder_type = self.db_folder.get_info(sensitive_info=True)["type"]
        if db_folder_type == "Filesystem" and not is_container_exec():
            return False
        else:
            return True

    def get_db_path(self) -> Path:
        """Get the absolute or relative path to the Kuzu database file.

        Returns:
            Path: The path to 'db.kz'.
        """
        relative_path_to_db = Path(BUILT_GRAPHS_PREFIX) / self.snapshot_id

        if self.is_remote:
            return relative_path_to_db / DB_KUZU_FILENAME
        else:
            return Path(self.db_folder.get_path()) / relative_path_to_db / DB_KUZU_FILENAME


class Neo4jConnectionInfo(AbstractConnectionInfo):
    """Connection information for a Neo4j database.

    Attributes:
        database_name (str): The name of the database in Neo4j.
        preset (Neo4jPresetConfig): Connection credentials and URI.
        db_folder (Optional[Folder]): The folder where metadata is stored, 
            if managed by DSS.
    """

    def __init__(self, snapshot_id: str, database_name: str, preset: Neo4jPresetConfig, db_folder: Optional[Folder] = None) -> None:
        """Initializes Neo4j connection info.

        Args:
            snapshot_id (str): Unique identifier for the snapshot.
            database_name (str): Name of the database in Neo4j.
            preset (Neo4jPresetConfig): Neo4j connection configuration.
            db_folder (Optional[Folder]): The DSS folder managing this connection, if any.
        """
        self.database_name = database_name
        self.snapshot_id = snapshot_id
        self.preset = preset
        self._db_folder = db_folder

    @property
    def type(self) -> DbType:
        """The type of the database ('neo4j')."""
        return "neo4j"

    @property
    def uid(self) -> str:
        """The unique identifier for this connection (snapshot ID)."""
        return self.snapshot_id
    
    @property
    def db_folder(self) -> Folder | None:
        """The DSS folder where the Neo4j metadata is stored, if managed."""
        return self._db_folder
    
    @property
    def db_name(self) -> str:
        """The name of the database in Neo4j."""
        return self.database_name

    @property
    def is_external(self) -> bool:
        """True if the database is not managed by a DSS graph folder."""
        return self.db_folder is None


def get_database_type_from_snapshot_folder(folder: Folder, snapshot_id: str) -> DbType:
    """Read the database type for a snapshot directly from the folder metadata.

    Defaults to Kuzu when `database.json` is missing, which matches the historical
    behavior used for older snapshots.
    """
    database_json_path = f"/{BUILT_GRAPHS_PREFIX}/{snapshot_id}/{DATABASE_JSON_FILENAME}"
    try:
        database_json = folder.read_json(database_json_path)
        database_type = database_json.get("type")
    except Exception as e:
        logger.warning(f"Error reading database.json for snapshot {snapshot_id}: {e}. Defaulting to kuzu.")
        return "kuzu"

    if not database_type or database_type == "kuzu":
        return "kuzu"
    if database_type == "neo4j":
        return "neo4j"
    raise ValueError(f"Unknown database type '{database_type}' for snapshot {snapshot_id}.")


def get_connection_info_from_folder(
    folder: Folder, snapshot_id: str, preset: Optional[Neo4jPresetConfig] = None
) -> AbstractConnectionInfo:
    """Factory function to create ConnectionInfo from a DSS folder.

    Reads 'database.json' to determine the database type and parameters.

    Args:
        folder (Folder): The DSS folder to scan.
        snapshot_id (str): The ID of the snapshot to look up.
        preset (Optional[Neo4jPresetConfig]): Neo4j credentials if applicable.

    Returns:
        AbstractConnectionInfo: Either KuzuConnectionInfo or Neo4jConnectionInfo.

    Raises:
        ValueError: If the database type is unknown or configuration is missing.
    """
    database_json_path = f"/{BUILT_GRAPHS_PREFIX}/{snapshot_id}/{DATABASE_JSON_FILENAME}"
    database_type = get_database_type_from_snapshot_folder(folder, snapshot_id)
    database_json = None
    if database_type == "neo4j":
        database_json = folder.read_json(database_json_path)

    if database_type == "kuzu":
        return KuzuConnectionInfo(folder, snapshot_id)
    elif database_type == "neo4j":
        if not preset:
            raise ValueError(f"No Neo4j preset found for webapp")
        if not database_json:
            raise ValueError(f"No database.json found for snapshot {snapshot_id}")

        database_name = database_json["params"]["databaseName"]
        return Neo4jConnectionInfo(snapshot_id, database_name, preset, folder)
    else:
        raise ValueError(f"Unknown database type '{database_type}' for snapshot {snapshot_id}.")
