from __future__ import annotations

import logging
from enum import Enum
from functools import lru_cache
from typing import Any, Callable, Generator, Literal

import pandas as pd

from graph.connections.connection_info import AbstractConnectionInfo
from graph.connections.db_connection import DbConnection
from graph.connections.db_handler import DbHandler

GraphAnalyticsAlgorithmName = Literal["page_rank", "graph_features", "graph_clustering"]
GraphAnalyticsExecutor = Callable[..., Generator[pd.DataFrame, Any, None]]
logger = logging.getLogger(__name__)


class UnsupportedGraphAnalyticsBackendError(ValueError):
    """Raised when an algorithm is not supported for a given backend connection."""


class NativeGraphAnalyticsUnavailableError(ValueError):
    """Raised when native execution is required but unavailable."""


class InMemoryExecutionOutOfMemoryError(MemoryError):
    """Raised when explicit in-memory graph analytics execution exceeds available memory."""


class GraphAnalyticsExecutionLocation(str, Enum):
    NATIVE = "native"
    IN_MEMORY = "in_memory"


class GraphAnalyticsSamplingMode(str, Enum):
    NONE = "none"
    TOP_N_ROWS = "top_n_rows"


def _rewrite_in_memory_execution_error(
    error: Exception,
    *,
    location: GraphAnalyticsExecutionLocation,
    sampling_mode: GraphAnalyticsSamplingMode,
) -> Exception:
    if location != GraphAnalyticsExecutionLocation.IN_MEMORY or not isinstance(error, MemoryError):
        return error

    if sampling_mode == GraphAnalyticsSamplingMode.NONE:
        return InMemoryExecutionOutOfMemoryError(
            "The selected graph is too large for in-memory full execution. "
            "Retry with sampling enabled to run the algorithm on a bounded sample."
        )

    return InMemoryExecutionOutOfMemoryError(
        "The selected graph sample is too large for in-memory execution. "
        "Reduce the sample limit per group and try again."
    )


class GraphAnalyticsCapability:
    """Execution capabilities available for an algorithm on a given backend."""

    def __init__(self, native: bool = False, in_memory_fallback: bool = False) -> None:
        self.native = native
        self.in_memory_fallback = in_memory_fallback

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, GraphAnalyticsCapability):
            return NotImplemented
        return (
            self.native == other.native
            and self.in_memory_fallback == other.in_memory_fallback
        )

    def supported_locations(self) -> tuple[str, ...]:
        locations: list[str] = []
        if self.native:
            locations.append("native")
        if self.in_memory_fallback:
            locations.append("in-memory")
        return tuple(locations)


class GraphAnalyticsBackendRoute:
    """Execution routes registered for an algorithm/backend pair."""

    def __init__(
        self,
        connection_type: type[DbConnection],
        native_executor: GraphAnalyticsExecutor | None = None,
        native_unavailable_errors: tuple[type[Exception], ...] = (),
        native_sampling_modes: tuple[GraphAnalyticsSamplingMode, ...] = (),
        in_memory_sampling_modes: tuple[GraphAnalyticsSamplingMode, ...] = (),
    ) -> None:
        self.connection_type = connection_type
        self.native_executor = native_executor
        self.native_unavailable_errors = native_unavailable_errors
        self.native_sampling_modes = native_sampling_modes
        self.in_memory_sampling_modes = in_memory_sampling_modes


class GraphAnalyticsExecution:
    """Execution metadata plus a managed batch generator for a graph analytics run."""

    def __init__(
        self,
        connection_info: AbstractConnectionInfo,
        batches: Generator[pd.DataFrame, Any, None],
        cleanup_cb: Callable[[], None],
    ) -> None:
        self.connection_info = connection_info
        self.batches = batches
        self._cleanup_cb = cleanup_cb
        self._closed = False

    def __enter__(self) -> GraphAnalyticsExecution:
        return self

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

    def close(self) -> None:
        if self._closed:
            return

        try:
            self.batches.close()
        finally:
            self._cleanup_cb()
            self._closed = True


class GraphAnalyticsBackendService:
    """Dispatch graph analytics algorithms to the matching backend implementation."""

    def __init__(self) -> None:
        self._routes: dict[GraphAnalyticsAlgorithmName, list[GraphAnalyticsBackendRoute]] = {}
        self._fallback_executors: dict[GraphAnalyticsAlgorithmName, GraphAnalyticsExecutor] = {}

    def register(
        self,
        algorithm_name: GraphAnalyticsAlgorithmName,
        connection_type: type[DbConnection],
        native_executor: GraphAnalyticsExecutor | None = None,
        *,
        fallback_executor: GraphAnalyticsExecutor | None = None,
        native_unavailable_errors: tuple[type[Exception], ...] = (),
        native_sampling_modes: tuple[GraphAnalyticsSamplingMode, ...] = (),
        in_memory_sampling_modes: tuple[GraphAnalyticsSamplingMode, ...] = (),
    ) -> None:
        if native_executor is None and fallback_executor is None:
            raise ValueError("At least one executor must be registered.")

        if fallback_executor is not None:
            registered_fallback = self._fallback_executors.get(algorithm_name)
            if registered_fallback is None:
                self._fallback_executors[algorithm_name] = fallback_executor
            elif registered_fallback is not fallback_executor:
                raise ValueError(f"Fallback executor already registered for algorithm: {algorithm_name}")

        self._routes.setdefault(algorithm_name, []).append(
            GraphAnalyticsBackendRoute(
                connection_type=connection_type,
                native_executor=native_executor,
                native_unavailable_errors=native_unavailable_errors,
                native_sampling_modes=native_sampling_modes,
                in_memory_sampling_modes=in_memory_sampling_modes,
            )
        )

    def _get_route(
        self,
        algorithm_name: GraphAnalyticsAlgorithmName,
        connection_type: type[DbConnection] | DbConnection,
    ) -> GraphAnalyticsBackendRoute | None:
        routes = self._routes.get(algorithm_name)
        if not routes:
            return None

        resolved_connection_type = connection_type if isinstance(connection_type, type) else type(connection_type)
        for route in routes:
            if issubclass(resolved_connection_type, route.connection_type):
                return route

        return None

    def get_capability(
        self,
        algorithm_name: GraphAnalyticsAlgorithmName,
        connection_type: type[DbConnection] | DbConnection,
    ) -> GraphAnalyticsCapability:
        route = self._get_route(algorithm_name, connection_type)
        if route is None:
            return GraphAnalyticsCapability()
        return GraphAnalyticsCapability(
            native=route.native_executor is not None,
            in_memory_fallback=algorithm_name in self._fallback_executors,
        )

    def list_algorithms(self) -> tuple[GraphAnalyticsAlgorithmName, ...]:
        return tuple(sorted(self._routes))

    def list_capabilities(
        self,
        connection_type: type[DbConnection] | DbConnection,
    ) -> dict[GraphAnalyticsAlgorithmName, GraphAnalyticsCapability]:
        return {
            algorithm_name: self.get_capability(algorithm_name, connection_type)
            for algorithm_name in self.list_algorithms()
        }

    def get_supported_sampling_modes(
        self,
        algorithm_name: GraphAnalyticsAlgorithmName,
        connection_type: type[DbConnection] | DbConnection,
        location: GraphAnalyticsExecutionLocation | str,
    ) -> tuple[GraphAnalyticsSamplingMode, ...]:
        route = self._get_route(algorithm_name, connection_type)
        if route is None:
            return ()

        resolved_location = GraphAnalyticsExecutionLocation(location)
        if resolved_location == GraphAnalyticsExecutionLocation.NATIVE:
            return route.native_sampling_modes
        return route.in_memory_sampling_modes

    def execute(
        self,
        algorithm_name: GraphAnalyticsAlgorithmName,
        conn: DbConnection,
        *args: Any,
        batch_size: int = 10000,
        location: GraphAnalyticsExecutionLocation = GraphAnalyticsExecutionLocation.NATIVE,
        sampling_mode: GraphAnalyticsSamplingMode | None = None,
        sampling_limit: int | None = None,
    ) -> Generator[pd.DataFrame, Any, None]:
        resolved_sampling_mode = sampling_mode or GraphAnalyticsSamplingMode.NONE
        _validate_sampling_inputs(
            resolved_sampling_mode,
            sampling_limit=sampling_limit,
        )

        native_kwargs: dict[str, object] = {"batch_size": batch_size}
        fallback_kwargs: dict[str, object] = {"batch_size": batch_size}
        if resolved_sampling_mode != GraphAnalyticsSamplingMode.NONE:
            fallback_kwargs["sampling_mode"] = resolved_sampling_mode
            fallback_kwargs["sampling_limit"] = sampling_limit
            native_kwargs["sampling_mode"] = resolved_sampling_mode
            native_kwargs["sampling_limit"] = sampling_limit

        route = self._get_route(algorithm_name, conn)
        if route is None:
            if algorithm_name not in self._routes:
                raise ValueError(f"Unsupported graph algorithm: {algorithm_name}")
            raise UnsupportedGraphAnalyticsBackendError(f"Unsupported connection type: {type(conn).__name__}")
        fallback_executor = self._fallback_executors.get(algorithm_name)

        def batches() -> Generator[pd.DataFrame, Any, None]:
            if location == GraphAnalyticsExecutionLocation.IN_MEMORY:
                if fallback_executor is None:
                    raise UnsupportedGraphAnalyticsBackendError(
                        f"In-memory execution is not supported for connection type: {type(conn).__name__}"
                    )
                if (
                    resolved_sampling_mode != GraphAnalyticsSamplingMode.NONE
                    and resolved_sampling_mode not in route.in_memory_sampling_modes
                ):
                    raise UnsupportedGraphAnalyticsBackendError(
                        f"Sampling mode '{resolved_sampling_mode.value}' is not supported for in-memory execution on {type(conn).__name__}."
                    )
                execution_label = _describe_execution(location, resolved_sampling_mode)
                logger.info(
                    "Executing graph analytics algorithm '%s' in %s mode on %s.",
                    algorithm_name,
                    execution_label,
                    type(conn).__name__,
                )
                try:
                    yield from fallback_executor(conn, *args, **fallback_kwargs)
                except Exception as exc:
                    rewritten_error = _rewrite_in_memory_execution_error(
                        exc,
                        location=location,
                        sampling_mode=resolved_sampling_mode,
                    )
                    if rewritten_error is exc:
                        raise
                    raise rewritten_error from exc
                return

            if route.native_executor is None:
                raise NativeGraphAnalyticsUnavailableError(
                    f"Native execution is not supported for connection type: {type(conn).__name__}"
                )
            if (
                resolved_sampling_mode != GraphAnalyticsSamplingMode.NONE
                and resolved_sampling_mode not in route.native_sampling_modes
            ):
                raise UnsupportedGraphAnalyticsBackendError(
                    f"Sampling mode '{resolved_sampling_mode.value}' is not supported for native execution on {type(conn).__name__}."
                )

            try:
                logger.info(
                    "Executing graph analytics algorithm '%s' in %s mode on %s.",
                    algorithm_name,
                    _describe_execution(location, resolved_sampling_mode),
                    type(conn).__name__,
                )
                yield from route.native_executor(conn, *args, **native_kwargs)
            except route.native_unavailable_errors as exc:
                detail = f": {exc}" if str(exc) else ""
                raise NativeGraphAnalyticsUnavailableError(
                    f"Native execution is unavailable for connection type: {type(conn).__name__}{detail}"
                ) from exc

        return batches()


def execute_graph_algorithm(
    algorithm_name: GraphAnalyticsAlgorithmName,
    connection_info: AbstractConnectionInfo,
    db_handler: DbHandler,
    *args: Any,
    readonly: bool = False,
    **kwargs: Any,
) -> GraphAnalyticsExecution:
    def batches() -> Generator[pd.DataFrame, Any, None]:
        with db_handler.get_new_db_instance(readonly=readonly) as db_instance:
            with db_instance.get_new_conn() as connection:
                yield from get_graph_analytics_backend_service().execute(algorithm_name, connection, *args, **kwargs)

    return GraphAnalyticsExecution(connection_info=connection_info, batches=batches(), cleanup_cb=db_handler.cleanup_handler)


def build_recipe_execution_note(
    location: GraphAnalyticsExecutionLocation,
    sampling_mode: GraphAnalyticsSamplingMode,
) -> str:
    if location == GraphAnalyticsExecutionLocation.NATIVE:
        if sampling_mode == GraphAnalyticsSamplingMode.NONE:
            return "Execution engine is In database: the full graph is processed directly in the graph database."
        return (
            "Execution engine is In database with sampling enabled. "
            "The recipe will fail if the selected backend does not support this sampling mode."
        )
    if sampling_mode == GraphAnalyticsSamplingMode.NONE:
        return "Execution engine is DSS: the full selected graph is loaded in DSS backend memory."
    return (
        f"Execution engine is DSS with {sampling_mode.value} sampling: "
        "only a sampled subgraph will be processed."
    )


def _validate_sampling_inputs(
    sampling_mode: GraphAnalyticsSamplingMode,
    *,
    sampling_limit: int | None,
) -> None:
    if sampling_mode == GraphAnalyticsSamplingMode.TOP_N_ROWS:
        if sampling_limit is None or sampling_limit <= 0:
            raise ValueError("sampling_limit must be provided and greater than 0 when sampling is enabled")
        return


def _describe_execution(
    location: GraphAnalyticsExecutionLocation,
    sampling_mode: GraphAnalyticsSamplingMode,
) -> str:
    if location == GraphAnalyticsExecutionLocation.NATIVE:
        if sampling_mode == GraphAnalyticsSamplingMode.NONE:
            return "native"
        return f"native with {sampling_mode.value} sampling"
    if sampling_mode == GraphAnalyticsSamplingMode.NONE:
        return "in-memory full"
    return f"in-memory with {sampling_mode.value} sampling"


@lru_cache(maxsize=1)
def get_graph_analytics_backend_service() -> GraphAnalyticsBackendService:
    from graph.connections.kuzu.db_connection import KuzuConnection
    from graph.connections.neo4j.db_connection import Neo4jConnection
    from graph.exploration.algorithms.in_memory.graph_clustering import (
        compute_graph_clustering as compute_graph_clustering_in_memory,
    )
    from graph.exploration.algorithms.in_memory.graph_features import (
        compute_graph_features as compute_graph_features_in_memory,
    )
    from graph.exploration.algorithms.in_memory.page_rank import compute_page_rank as compute_page_rank_in_memory
    from graph.exploration.algorithms.kuzu.page_rank import compute_page_rank as compute_page_rank_kuzu
    from graph.exploration.algorithms.neo4j.page_rank import Neo4jGDSNotAvailableError
    from graph.exploration.algorithms.neo4j.page_rank import compute_page_rank as compute_page_rank_neo4j

    service = GraphAnalyticsBackendService()
    service.register(
        "page_rank",
        KuzuConnection,
        compute_page_rank_kuzu,
        fallback_executor=compute_page_rank_in_memory,
        in_memory_sampling_modes=(
            GraphAnalyticsSamplingMode.TOP_N_ROWS,
        ),
    )
    service.register(
        "page_rank",
        Neo4jConnection,
        compute_page_rank_neo4j,
        fallback_executor=compute_page_rank_in_memory,
        native_unavailable_errors=(Neo4jGDSNotAvailableError,),
        in_memory_sampling_modes=(
            GraphAnalyticsSamplingMode.TOP_N_ROWS,
        ),
    )
    service.register(
        "graph_features",
        KuzuConnection,
        fallback_executor=compute_graph_features_in_memory,
        in_memory_sampling_modes=(
            GraphAnalyticsSamplingMode.TOP_N_ROWS,
        ),
    )
    service.register(
        "graph_clustering",
        KuzuConnection,
        fallback_executor=compute_graph_clustering_in_memory,
        in_memory_sampling_modes=(
            GraphAnalyticsSamplingMode.TOP_N_ROWS,
        ),
    )
    service.register(
        "graph_features",
        Neo4jConnection,
        fallback_executor=compute_graph_features_in_memory,
        in_memory_sampling_modes=(
            GraphAnalyticsSamplingMode.TOP_N_ROWS,
        ),
    )
    service.register(
        "graph_clustering",
        Neo4jConnection,
        fallback_executor=compute_graph_clustering_in_memory,
        in_memory_sampling_modes=(
            GraphAnalyticsSamplingMode.TOP_N_ROWS,
        ),
    )
    return service
