from __future__ import annotations

from collections.abc import Hashable
from dataclasses import dataclass
from typing import Any

import pandas as pd

from graph.connections.db_connection import DbConnection
from graph.connections.kuzu.db_connection import KuzuConnection
from graph.connections.neo4j.db_connection import Neo4jConnection
from graph.exploration.algorithms.backend_service import GraphAnalyticsSamplingMode
from graph.exploration.algorithms.kuzu.common import get_node_groups_info as get_kuzu_node_groups_info
from graph.exploration.algorithms.neo4j.common import get_node_groups_info as get_neo4j_node_groups_info
from graph.exploration.algorithms.projected_graph import ProjectedGraph


@dataclass(frozen=True)
class LoadedNode:
    key: Hashable
    node_group: str
    node_id: Any
    raw: dict[str, Any]


@dataclass(frozen=True)
class LoadedEdge:
    source_key: Hashable
    target_key: Hashable
    source_group: str
    source_id: Any
    target_group: str
    target_id: Any
    edge_group: str
    weight: float | None = None


@dataclass(frozen=True)
class LoadedProjectedGraph:
    nodes: dict[Hashable, LoadedNode]
    edges: list[LoadedEdge]


def quote_identifier(identifier: str) -> str:
    return f"`{identifier.replace('`', '``')}`"


def _resolve_internal_id(node: dict[str, Any]) -> Any:
    element_id = node.get("_element_id")
    if element_id is not None:
        return element_id
    return node.get("_id")


def extract_internal_node_key(node: dict[str, Any]) -> Hashable:
    internal_id = _resolve_internal_id(node)
    if isinstance(internal_id, dict):
        return (str(internal_id.get("table")), str(internal_id.get("offset")))
    if isinstance(internal_id, Hashable):
        return internal_id

    node_label = node.get("_label")
    primary_key_name = node.get("_dku_reserved_user_defined_pk_prop")
    if node_label and isinstance(primary_key_name, str) and primary_key_name in node:
        return (node_label, node[primary_key_name])

    raise ValueError(f"Unable to infer a stable node key for node: {node}")


def extract_output_node_id(node: dict[str, Any], primary_key_name: str | None) -> Any:
    if primary_key_name == "_internal_id":
        return _resolve_internal_id(node)

    if primary_key_name and primary_key_name in node:
        return node[primary_key_name]

    fallback_key_name = node.get("_dku_reserved_user_defined_pk_prop")
    if isinstance(fallback_key_name, str) and fallback_key_name in node:
        return node[fallback_key_name]

    if "id" in node:
        return node["id"]

    return _resolve_internal_id(node)


def _get_node_groups_info(conn: DbConnection, projected_graph: ProjectedGraph) -> dict[str, str]:
    if isinstance(conn, KuzuConnection):
        return get_kuzu_node_groups_info(conn)
    if isinstance(conn, Neo4jConnection):
        return get_neo4j_node_groups_info(conn, projected_graph)
    raise ValueError(f"Unsupported connection type for projected graph loading: {type(conn).__name__}")


def _load_nodes_from_df(
    raw_nodes: dict[Hashable, dict[str, Any]],
    df: pd.DataFrame,
) -> None:
    for node in df["n"]:
        if isinstance(node, dict):
            raw_nodes[extract_internal_node_key(node)] = node


def _load_edges_from_df(
    raw_nodes: dict[Hashable, dict[str, Any]],
    raw_edges: list[tuple[Hashable, Hashable, str, float | None]],
    df: pd.DataFrame,
    *,
    allowed_node_keys: set[Hashable] | None = None,
    weight_property: str | None = None,
) -> None:
    for source, edge, target in zip(df["s"], df["r"], df["t"]):
        if not isinstance(source, dict) or not isinstance(edge, dict) or not isinstance(target, dict):
            continue

        source_key = extract_internal_node_key(source)
        target_key = extract_internal_node_key(target)
        if allowed_node_keys is not None and (source_key not in allowed_node_keys or target_key not in allowed_node_keys):
            continue

        raw_nodes[source_key] = source
        raw_nodes[target_key] = target
        raw_edges.append(
            (
                source_key,
                target_key,
                str(edge.get("_label", "")),
                _extract_edge_weight_value(edge, weight_property),
            )
        )


def _extract_edge_weight_value(edge_data: dict[str, Any], weight_property: str | None) -> float | None:
    if not weight_property:
        return None

    if weight_property not in edge_data:
        raise ValueError(f"Selected weight property '{weight_property}' does not exist on all selected edges.")

    value = edge_data[weight_property]
    if isinstance(value, bool):
        return float(value)
    if isinstance(value, (int, float)):
        return float(value)

    raise ValueError(f"Selected weight property '{weight_property}' must contain numeric values.")


def _quote_cypher_value(value: Any) -> str:
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, (int, float)):
        return str(value)
    escaped = str(value).replace("\\", "\\\\").replace("'", "\\'")
    return f"'{escaped}'"


def _build_in_list(values: list[Any]) -> str:
    return "[" + ", ".join(_quote_cypher_value(value) for value in values) + "]"


def _get_kuzu_edge_endpoints(conn: KuzuConnection, edge_group: str) -> tuple[str, str]:
    connection_df = conn.execute(f'CALL SHOW_CONNECTION("{edge_group}") RETURN *;').get_as_df()  # type: ignore
    source_group = str(connection_df.iloc[0]["source table name"])
    target_group = str(connection_df.iloc[0]["destination table name"])
    return source_group, target_group


def _sampled_neo4j_element_ids(raw_nodes: dict[Hashable, dict[str, Any]]) -> list[str]:
    return sorted(
        {
            node_id
            for node in raw_nodes.values()
            if isinstance((node_id := node.get("_id")), str)
        }
    )


def _sampled_kuzu_offsets(raw_nodes: dict[Hashable, dict[str, Any]], node_group: str) -> list[int]:
    offsets: set[int] = set()
    for node in raw_nodes.values():
        if node.get("_label") != node_group:
            continue
        internal_id = node.get("_id")
        if isinstance(internal_id, dict):
            offset = internal_id.get("offset")
            if isinstance(offset, int):
                offsets.add(offset)
    return sorted(offsets)


def _build_sampled_edge_query(
    conn: DbConnection,
    edge_group: str,
    raw_nodes: dict[Hashable, dict[str, Any]],
    sample_limit: int | None,
) -> str | None:
    limit_clause = f" LIMIT {sample_limit}" if sample_limit is not None else ""

    if isinstance(conn, Neo4jConnection):
        sampled_element_ids = _sampled_neo4j_element_ids(raw_nodes)
        if not sampled_element_ids:
            return None
        sampled_ids_list = _build_in_list(sampled_element_ids)
        return (
            f"MATCH (s)-[r:{quote_identifier(edge_group)}]->(t) "
            f"WHERE elementId(s) IN {sampled_ids_list} AND elementId(t) IN {sampled_ids_list} "
            f"RETURN s, r, t ORDER BY id(s), id(t), id(r){limit_clause}"
        )

    if isinstance(conn, KuzuConnection):
        source_group, target_group = _get_kuzu_edge_endpoints(conn, edge_group)
        source_offsets = _sampled_kuzu_offsets(raw_nodes, source_group)
        target_offsets = _sampled_kuzu_offsets(raw_nodes, target_group)
        if not source_offsets or not target_offsets:
            return None
        return (
            f"MATCH (s)-[r:{quote_identifier(edge_group)}]->(t) "
            f"WHERE OFFSET(ID(s)) IN {_build_in_list(source_offsets)} "
            f"AND OFFSET(ID(t)) IN {_build_in_list(target_offsets)} "
            f"RETURN s, r, t ORDER BY OFFSET(ID(s)), OFFSET(ID(t)), OFFSET(ID(r)){limit_clause}"
        )

    return None


def _build_node_query_suffix(
    conn: DbConnection,
    sampling_mode: GraphAnalyticsSamplingMode,
    sampling_limit: int | None,
) -> str:
    if isinstance(conn, Neo4jConnection):
        order_clause = " ORDER BY id(n)"
    elif isinstance(conn, KuzuConnection):
        order_clause = " ORDER BY OFFSET(ID(n))"
    else:
        order_clause = ""

    if sampling_mode == GraphAnalyticsSamplingMode.NONE:
        return order_clause
    if sampling_mode == GraphAnalyticsSamplingMode.TOP_N_ROWS:
        if sampling_limit is None or sampling_limit <= 0:
            raise ValueError("sampling_limit is required and must be greater than 0 for top_n_rows sampling.")
        return f"{order_clause} LIMIT {sampling_limit}"
    raise ValueError(f"Unsupported in-memory sampling mode '{sampling_mode.value}'.")


def _build_full_edge_query(conn: DbConnection, edge_group: str) -> str:
    order_clause = ""
    if isinstance(conn, Neo4jConnection):
        order_clause = " ORDER BY id(s), id(t), id(r)"
    elif isinstance(conn, KuzuConnection):
        order_clause = " ORDER BY OFFSET(ID(s)), OFFSET(ID(t)), OFFSET(ID(r))"
    return f"MATCH (s)-[r:{quote_identifier(edge_group)}]->(t) RETURN s, r, t{order_clause}"


def load_projected_graph(
    conn: DbConnection,
    projected_graph: ProjectedGraph,
    *,
    sampling_mode: GraphAnalyticsSamplingMode = GraphAnalyticsSamplingMode.NONE,
    sampling_limit: int | None = None,
    weight_property: str | None = None,
) -> LoadedProjectedGraph:
    resolved_sampling_mode = sampling_mode
    if resolved_sampling_mode == GraphAnalyticsSamplingMode.NONE and sampling_limit is not None:
        resolved_sampling_mode = GraphAnalyticsSamplingMode.TOP_N_ROWS

    node_groups_info = _get_node_groups_info(conn, projected_graph)
    raw_nodes: dict[Hashable, dict[str, Any]] = {}
    raw_edges: list[tuple[Hashable, Hashable, str, float | None]] = []
    allowed_node_keys: set[Hashable] | None = set() if resolved_sampling_mode != GraphAnalyticsSamplingMode.NONE else None

    for node_group in sorted(projected_graph.node_groups):
        query_suffix = _build_node_query_suffix(conn, resolved_sampling_mode, sampling_limit)
        result = conn.execute(f"MATCH (n:{quote_identifier(node_group)}) RETURN n{query_suffix}")
        try:
            _load_nodes_from_df(raw_nodes, result.get_as_df())
            if allowed_node_keys is not None:
                allowed_node_keys.update(raw_nodes)
        finally:
            result.close()

    for edge_group in sorted(projected_graph.edge_groups):
        if resolved_sampling_mode != GraphAnalyticsSamplingMode.NONE:
            query = _build_sampled_edge_query(conn, edge_group, raw_nodes, sampling_limit)
            if query is None:
                continue
        else:
            query = _build_full_edge_query(conn, edge_group)

        result = conn.execute(query)
        try:
            _load_edges_from_df(
                raw_nodes,
                raw_edges,
                result.get_as_df(),
                allowed_node_keys=allowed_node_keys,
                weight_property=weight_property,
            )
        finally:
            result.close()

    nodes = {
        node_key: LoadedNode(
            key=node_key,
            node_group=str(node_data.get("_label")),
            node_id=extract_output_node_id(
                node_data,
                node_groups_info.get(str(node_data.get("_label"))),
            ),
            raw=node_data,
        )
        for node_key, node_data in raw_nodes.items()
        if node_data.get("_label") in projected_graph.node_groups
    }

    edges = [
        LoadedEdge(
            source_key=source_key,
            target_key=target_key,
            source_group=nodes[source_key].node_group,
            source_id=nodes[source_key].node_id,
            target_group=nodes[target_key].node_group,
            target_id=nodes[target_key].node_id,
            edge_group=edge_group,
            weight=edge_weight,
        )
        for source_key, target_key, edge_group, edge_weight in raw_edges
        if source_key in nodes and target_key in nodes
    ]

    return LoadedProjectedGraph(nodes=nodes, edges=edges)
