from __future__ import annotations

from contextlib import ExitStack
from dataclasses import dataclass
from itertools import chain
from typing import Any, Callable, Iterator, cast

import pandas as pd
from dataiku import Dataset, Folder

from backends.explorer.discovery.connection_info import AbstractConnectionInfo
from config_utils import strtobool
from graph.connections.db_connection import DbConnection
from graph.connections.db_handler import DbHandler
from graph.connections.kuzu.db_connection import KuzuConnection
from graph.connections.neo4j.db_connection import Neo4jConnection
from graph.connections.neo4j.preset import Neo4jPresetConfig
from graph.exploration.algorithms.backend_service import (
    GraphAnalyticsAlgorithmName,
    GraphAnalyticsCapability,
    GraphAnalyticsExecutionLocation,
    GraphAnalyticsSamplingMode,
    get_graph_analytics_backend_service,
)
from graph.exploration.algorithms.graph_algorithm_registry import (
    GraphAlgorithmDefinition,
    default_algorithm_names,
)
from recipes.explore_graph.utils import get_db_handler, get_single_connection_info
from solutions.graph.models import DbType


@dataclass(frozen=True)
class GraphAnalyticsRecipeBackendConfig:
    neo4j_configuration: Neo4jPresetConfig | None
    database_name: str | None
    database_type: str | None
    location: GraphAnalyticsExecutionLocation
    sampling_mode: GraphAnalyticsSamplingMode
    sampling_limit: int | None

    @staticmethod
    def get_effective_database_type(database_type: str | None) -> DbType | None:
        if database_type == "unmanaged_neo4j":
            return "neo4j"
        if database_type in ("kuzu", "neo4j"):
            return cast(DbType, database_type)
        return None

    @classmethod
    def from_recipe_config(
        cls,
        config: dict[str, Any],
        *,
        default_location: GraphAnalyticsExecutionLocation,
    ) -> GraphAnalyticsRecipeBackendConfig:
        neo4j_configuration = config.get("neo4j_preset")
        if neo4j_configuration:
            neo4j_configuration = Neo4jPresetConfig.from_dict(neo4j_configuration)

        location = _resolve_recipe_location(config.get("location"), default_location)
        sampling_enabled = parse_recipe_bool(config.get("sampling_enabled"), default=False)
        if sampling_enabled:
            sampling_mode = GraphAnalyticsSamplingMode(
                config.get("sampling_mode") or GraphAnalyticsSamplingMode.TOP_N_ROWS
            )
        else:
            sampling_mode = GraphAnalyticsSamplingMode.NONE

        sampling_limit = None
        if sampling_mode == GraphAnalyticsSamplingMode.TOP_N_ROWS:
            sampling_limit = int(config.get("sampling_limit_per_group", 1000))
            if sampling_limit <= 0:
                raise ValueError("sampling_limit_per_group must be greater than 0 when sampling is enabled")

        return cls(
            neo4j_configuration=neo4j_configuration,
            database_name=config.get("database_name"),
            database_type=config.get("database_type"),
            location=location,
            sampling_mode=sampling_mode,
            sampling_limit=sampling_limit,
        )

    def resolve_backend(
        self,
        *,
        db_folder: Folder | None,
    ) -> tuple[DbType | None, AbstractConnectionInfo, DbHandler]:
        effective_database_type = self.get_effective_database_type(self.database_type)
        if self.database_type is not None and effective_database_type is None:
            raise ValueError(f"Unsupported graph analytics database type: {self.database_type}")

        conn_info = get_single_connection_info(
            db_folder=db_folder,
            neo4j_configuration=self.neo4j_configuration,
            database_name=self.database_name,
            database_type=effective_database_type,
        )
        db_handler = get_db_handler(
            db_folder=db_folder,
            neo4j_configuration=self.neo4j_configuration,
            database_name=self.database_name,
            database_type=effective_database_type,
            snapshot_id=conn_info.uid,
        )
        return effective_database_type, conn_info, db_handler


def write_batches_to_dataset(
    output_ds: Dataset,
    batches: Iterator[pd.DataFrame],
    *,
    empty_frame_factory: Callable[[], pd.DataFrame] | None = None,
    empty_error_factory: Callable[[], Exception] | None = None,
    on_batch_write: Callable[[pd.DataFrame], None] | None = None,
) -> None:
    first_batch = next(batches, None)
    if first_batch is None:
        if empty_frame_factory is not None:
            first_batch = empty_frame_factory()
        elif empty_error_factory is not None:
            raise empty_error_factory()
        else:
            return

    with ExitStack() as exit_stack:
        output_ds.write_schema_from_dataframe(first_batch, drop_and_create=True)
        output_writer = exit_stack.enter_context(output_ds.get_writer())

        for df in chain([first_batch], batches):
            if df.empty:
                continue
            if on_batch_write is not None:
                on_batch_write(df)
            output_writer.write_dataframe(df)


def validate_graph_analytics_recipe_backend(
    algorithm_name: GraphAnalyticsAlgorithmName,
    *,
    effective_database_type: DbType | None,
    backend_config: GraphAnalyticsRecipeBackendConfig,
) -> None:
    connection_type = _resolve_graph_analytics_connection_type(effective_database_type)
    service = get_graph_analytics_backend_service()
    capability = service.get_capability(algorithm_name, connection_type)
    _validate_requested_location(
        algorithm_name,
        backend_name=effective_database_type or "unknown",
        capability=capability,
        location=backend_config.location,
    )
    if backend_config.sampling_mode == GraphAnalyticsSamplingMode.NONE:
        return

    supported_sampling_modes = service.get_supported_sampling_modes(
        algorithm_name,
        connection_type,
        backend_config.location,
    )
    if backend_config.sampling_mode not in supported_sampling_modes:
        raise ValueError(
            f"{_format_algorithm_name(algorithm_name)} does not support "
            f"{_format_sampling_mode_label(backend_config.sampling_mode)} sampling for "
            f"{_format_location_label(backend_config.location)} execution on {_format_backend_name(effective_database_type)}. "
            f"Supported sampling modes: {_format_supported_sampling_modes(supported_sampling_modes)}."
        )


def parse_recipe_bool(value: Any, *, default: bool = False) -> bool:
    if value is None:
        return default
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        try:
            return bool(strtobool(value.strip()))
        except ValueError as error:
            raise ValueError(f"Invalid boolean recipe value: {value!r}") from error
    return bool(value)


def parse_graph_algorithm_selection(
    config: dict[str, Any],
    *,
    definitions: tuple[GraphAlgorithmDefinition, ...],
    directed_graph: bool,
) -> tuple[str, ...]:
    algorithm_field_names = {definition.name for definition in definitions}
    has_algorithm_fields = any(name in config for name in algorithm_field_names)

    if not has_algorithm_fields:
        select_all = parse_recipe_bool(config.get("select_all_algorithms"), default=True)
        return default_algorithm_names(definitions, directed_graph=directed_graph) if select_all else ()

    selected_algorithm_names: list[str] = []
    for definition in definitions:
        if directed_graph and not definition.supports_directed:
            continue
        if parse_recipe_bool(config.get(definition.name), default=False):
            selected_algorithm_names.append(definition.name)

    return tuple(selected_algorithm_names)


def _resolve_graph_analytics_connection_type(effective_database_type: DbType | None) -> type[DbConnection]:
    if effective_database_type == "kuzu":
        return KuzuConnection
    if effective_database_type == "neo4j":
        return Neo4jConnection
    raise ValueError(f"Unsupported graph analytics backend type: {effective_database_type!r}")


def _validate_requested_location(
    algorithm_name: GraphAnalyticsAlgorithmName,
    *,
    backend_name: str,
    capability: GraphAnalyticsCapability,
    location: GraphAnalyticsExecutionLocation,
) -> None:
    if location == GraphAnalyticsExecutionLocation.NATIVE and capability.native:
        return
    if location == GraphAnalyticsExecutionLocation.IN_MEMORY and capability.in_memory_fallback:
        return

    raise ValueError(
        f"{_format_algorithm_name(algorithm_name)} does not support "
        f"{_format_location_label(location)} execution on {_format_backend_name(backend_name)}. "
        f"Supported locations: {_format_supported_locations(capability.supported_locations())}."
    )


def _format_algorithm_name(algorithm_name: GraphAnalyticsAlgorithmName) -> str:
    return algorithm_name.replace("_", " ").title()


def _resolve_recipe_location(
    location: GraphAnalyticsExecutionLocation | str | None,
    default_location: GraphAnalyticsExecutionLocation,
) -> GraphAnalyticsExecutionLocation:
    if location is None:
        return GraphAnalyticsExecutionLocation(default_location)
    if isinstance(location, str) and not location.strip():
        return GraphAnalyticsExecutionLocation(default_location)
    return GraphAnalyticsExecutionLocation(location)


def _format_backend_name(backend_name: str | None) -> str:
    if backend_name == "neo4j":
        return "Neo4j"
    if backend_name == "kuzu":
        return "the built-in graph database"
    return str(backend_name)


def _format_location_label(location: GraphAnalyticsExecutionLocation) -> str:
    if location == GraphAnalyticsExecutionLocation.NATIVE:
        return "In database"
    return "DSS"


def _format_supported_locations(locations: tuple[str, ...]) -> str:
    if not locations:
        return "none"
    return ", ".join(_format_supported_location_label(location) for location in locations)


def _format_supported_location_label(location: str) -> str:
    if location == "native":
        return _format_location_label(GraphAnalyticsExecutionLocation.NATIVE)
    if location in ("in_memory", "in-memory"):
        return _format_location_label(GraphAnalyticsExecutionLocation.IN_MEMORY)
    return location


def _format_sampling_mode_label(sampling_mode: GraphAnalyticsSamplingMode) -> str:
    if sampling_mode == GraphAnalyticsSamplingMode.TOP_N_ROWS:
        return "Top N rows"
    return str(sampling_mode.value)


def _format_supported_sampling_modes(
    sampling_modes: tuple[GraphAnalyticsSamplingMode, ...],
) -> str:
    if not sampling_modes:
        return "none"
    return ", ".join(_format_sampling_mode_label(mode) for mode in sampling_modes)
