from __future__ import annotations

import logging
from typing import Any, Dict, cast

import kuzu
from networkx import MultiDiGraph
from pandas import DataFrame

from graph.connections.query_result import NodeData, QueryResult, RelationshipData, TableData
from graph.ingestion.dataiku.data_types import GraphDataikuDataType
from graph.ingestion.graph_builder import GROUP_ID_PROP_NAME
from graph.ingestion.kuzu.data_types import map_kuzu_to_dataiku_type

logger = logging.getLogger(__name__)


def get_as_networkx(query_result: kuzu.QueryResult) -> MultiDiGraph:
    """
    This code is heavily inspired from kuzu.QueryResult.get_as_networkx method and changed to fit our own usage.
    """
    query_result.check_for_query_result_close()
    properties_to_extract = query_result._get_properties_to_extract()

    query_result.reset_iterator()

    nodes: dict[tuple[int, int], dict[str, Any]] = {}
    rels = {}

    # This dictionary caches the primary key property name for each node table label.
    # For example: {'patient': '_dku_id_patients'}
    table_primary_key_dict: Dict[str, str] = {}

    def encode_node_id(node: dict[str, Any]) -> Any:
        node_label = node["_label"]
        group_id = node.get(GROUP_ID_PROP_NAME)

        # Group id may not exist for graphs created outside of VG editor/plugin.
        if group_id:
            # The '~' separator allows embedding the group ID within the node ID for later access.
            # The node's actual ID is its primary key value.
            # For example, a node {'_dku_grp_id': 'VodFnW', '_dku_id_patients': 2} results in the ID 'VodFnW~2'.
            # This is to ensure uniqueness of a node id across the whole graph. Two nodes in different group may have an identical identifier value.
            return f"{group_id}~{node[table_primary_key_dict[node_label]]!s}"
        else:
            return node[table_primary_key_dict[node_label]]

    def encode_rel_id(rel: dict[str, Any]) -> tuple[int, int]:
        return rel["_id"]["table"], rel["_id"]["offset"]

    def serialize_node(node: dict[str, Any]) -> tuple[str, NodeData]:
        """
        Serializes a Kuzu node dictionary into a standardized format.

        This function generates a unique node ID, extracts properties, and
        caches the primary key for the node's table label for efficiency.

        Args:
            node: A dictionary representing a Kuzu node.

        Returns:
            A tuple containing the unique node ID and a dictionary of node data.
        """
        node_label = node["_label"]

        # Find and cache the primary key for this node's label if we haven't seen it before.
        if node_label not in table_primary_key_dict:
            props = query_result.connection._get_node_property_names(node_label)  # type: ignore
            for prop_name in props:
                if props[prop_name]["is_primary_key"]:
                    table_primary_key_dict[node_label] = prop_name
                    break

        node_id = encode_node_id(node)
        del node["_id"]

        dku_grp_def_id = node.get("_dku_grp_def_id")
        assert dku_grp_def_id

        dku_reserved_user_defined_pk_prop = node.get("_dku_reserved_user_defined_pk_prop")
        assert dku_reserved_user_defined_pk_prop

        data = NodeData(
            id=node_id,
            _id=node_id,
            _label=node["_label"],
            _dku_grp_def_id=dku_grp_def_id,
            _dku_reserved_user_defined_pk_prop=dku_reserved_user_defined_pk_prop,
        )

        data.update(**dict(node))  # type: ignore

        return node_id, cast(NodeData, data)

    def serialize_relationship(rel: dict[str, Any]) -> tuple[str, str, RelationshipData] | None:
        """
        Serializes a Kuzu relationship dictionary into a standardized format.

        This function retrieves the source and destination nodes from the parent
        scope's `nodes` cache, generates their unique IDs, and constructs a
        standardized dictionary for the relationship.

        Args:
            rel: A dictionary representing a Kuzu relationship.

        Returns:
            A tuple containing the source node ID, destination node ID, and a
            dictionary of relationship data, or None if one endpoint node was
            not returned by the query.
        """
        _src = rel["_src"]
        _dst = rel["_dst"]
        src_node = nodes.get((_src["table"], _src["offset"]))
        dst_node = nodes.get((_dst["table"], _dst["offset"]))
        if src_node is None or dst_node is None:
            logger.debug(f"Skipping relationship {_src}-{_dst}: missing source or destination node.")
            return None

        src_id = encode_node_id(src_node)
        dst_id = encode_node_id(dst_node)

        dku_grp_def_id = rel.get("_dku_grp_def_id")
        assert dku_grp_def_id

        edge_data = RelationshipData(
            _src=_src,
            _dst=_dst,
            _label=rel["_label"],
            _dku_grp_def_id=dku_grp_def_id,
        )
        edge_data.update(**dict(rel))  # type: ignore

        return src_id, dst_id, cast(RelationshipData, edge_data)

    # De-duplicate nodes and rels
    while query_result.has_next():
        row = query_result.get_next()
        if not isinstance(row, list):
            raise Exception("Unexpected type when turning result in a networkx graph.")

        for i in properties_to_extract:
            # Skip empty nodes and rels, which may be returned by
            # OPTIONAL MATCH
            if row[i] is None or row[i] == {}:
                continue
            column_type, _ = properties_to_extract[i]
            if column_type == kuzu.Type.NODE.value:
                _id = row[i]["_id"]
                nodes[(_id["table"], _id["offset"])] = row[i]

            elif column_type == kuzu.Type.REL.value:
                rels[encode_rel_id(row[i])] = row[i]

            elif column_type == kuzu.Type.RECURSIVE_REL.value:
                for node in row[i]["_nodes"]:
                    _id = node["_id"]
                    nodes[(_id["table"], _id["offset"])] = node
                for rel in row[i]["_rels"]:
                    for key in list(rel.keys()):
                        if rel[key] is None:
                            del rel[key]
                    rels[encode_rel_id(rel)] = rel

    nx_graph: MultiDiGraph = MultiDiGraph()

    # Add nodes
    for node in nodes.values():
        node_id, node_data = serialize_node(node)
        nx_graph.add_node(node_id, **node_data)

    # Add rels
    for rel in rels.values():
        serialized_relationship = serialize_relationship(rel)
        if serialized_relationship is None:
            continue
        src_id, dst_id, edge_data = serialized_relationship
        nx_graph.add_edge(src_id, dst_id, **edge_data)

    return nx_graph


class KuzuQueryResult(QueryResult):
    def __init__(self, query_result: kuzu.QueryResult | list[kuzu.QueryResult]) -> None:
        self.__query_result = query_result

    def get_as_df(self) -> DataFrame:
        if isinstance(self.__query_result, list):
            raise RuntimeError("Multi query results are not supported.")

        return self.__query_result.get_as_df()

    def get_as_table(self) -> TableData:
        df = self.get_as_df()

        kuzu_schema: dict[str, str] = self.__query_result.get_schema()  # type: ignore
        dataiku_schema: dict[str, GraphDataikuDataType] = {
            kuzu_col: map_kuzu_to_dataiku_type(kuzu_type) for kuzu_col, kuzu_type in kuzu_schema.items()
        }

        return TableData(df=df, schema=dataiku_schema)

    def get_as_networkx(self) -> MultiDiGraph:
        if isinstance(self.__query_result, list):
            raise RuntimeError("Multi query results are not supported.")

        return get_as_networkx(self.__query_result)

    def close(self) -> None:
        # Still unsure of the origin:
        # if I dont explicitely close the QueryResult manually, it fails with an error regarding semaphores.
        # Might be due to the order in which the connection and the QueryResult are closed when holding variables go out-of-scope.
        if self.__query_result:
            if isinstance(self.__query_result, list):
                for q in self.__query_result:
                    try:
                        q.close()
                    except Exception:
                        logger.exception("Failed to close one of the query result.")
            else:
                try:
                    self.__query_result.close()
                except Exception:
                    logger.exception("Failed to close the query result.")
