from __future__ import annotations

import logging
from typing import List

from networkx import MultiDiGraph
from pandas import DataFrame

from backends.explorer.discovery.connection_info import AbstractConnectionInfo
from backends.explorer.metadata.metadata_builder import Colorizer, MetadataBuilder
from graph.connections.db_handler import DbHandler
from graph.connections.query_result import QueryResult
from solutions.graph.models import EdgeGroupDefinition, ExplorerMetadata, NodeGroupDefinition

logger = logging.getLogger(__name__)


class KuzuMetadataBuilder(MetadataBuilder):
    def __init__(self, db_handler: DbHandler, conn_info: AbstractConnectionInfo):
        super().__init__(db_handler, conn_info)

    def build(self) -> ExplorerMetadata:
        node_definitions: List[NodeGroupDefinition] = []
        edge_definitions: List[EdgeGroupDefinition] = []

        df = self.__execute_kuzu_query_df("CALL SHOW_TABLES() RETURN *;")
        df_dict = df.to_dict(orient="records")

        for col in df_dict:
            self._process_col(col, node_definitions, edge_definitions)

        return ExplorerMetadata(
            snapshot_id=self.conn_info.uid,
            db_type="kuzu",
            name="Unknown",
            comment="",
            epoch_ms=0,
            node_definitions=node_definitions,
            edge_definitions=edge_definitions,
            nodes_view=self.get_nodes_view(node_definitions),
            edges_view=self.get_edges_view(edge_definitions),
            cypher_queries=[],
        )

    def get_nodes_view(self, node_definitions: list) -> dict:
        node_ids = [node["node_id"] for node in node_definitions]
        return {
            node: {"color": Colorizer.string_to_beautiful_color(node), "size": "normal", "icon": ""}
            for node in node_ids
        }

    def get_edges_view(self, edge_definitions: list) -> dict:
        edge_ids = [edge["edge_id"] for edge in edge_definitions]
        return {edge: {"size": 1} for edge in edge_ids}

    def _process_col(
        self, col: dict, node_definitions: List[NodeGroupDefinition], edge_definitions: List[EdgeGroupDefinition]
    ) -> None:
        name = col["name"]

        if col["type"] == "NODE":
            # Query to find all the distinct nodes regarding _dku_grp_def_id. Each one is a node definition.
            query = f"MATCH (o:`{name}`) WITH o._dku_grp_def_id AS groupId, collect(o) AS objs RETURN objs[1] AS obj;"
            graph = self.__execute_kuzu_query_graph(query)
            for node, data in graph.nodes(data=True):
                node_id = node.split("~")[0]
                node_definitions.append(
                    NodeGroupDefinition(
                        definition_id=data["_dku_grp_def_id"],
                        node_group=data["_label"],
                        node_id=node_id,
                        source_dataset="",
                        primary_col=data["_dku_reserved_user_defined_pk_prop"],
                        label_col=data["_dku_reserved_user_defined_pk_prop"],
                        property_list=[k for k in data if not k.startswith("_")],
                        filters_association="and",
                        filters_stored=[],
                    )
                )
        else:
            # Query to find all the distinct edges regarding _dku_grp_def_id
            query = f"MATCH (o)-[e:`{name}`]->(d) WITH e._dku_grp_def_id AS groupId, collect({{source: o, rel: e, destination: d}}) AS paths WITH paths[1] AS firstPath RETURN firstPath.source AS o, firstPath.rel AS e, firstPath.destination AS d;"
            graph = self.__execute_kuzu_query_graph(query)
            for source, target, data in graph.edges(data=True):
                source_id = source.split("~")[0]
                target_id = target.split("~")[0]
                edge_definitions.append(
                    EdgeGroupDefinition(
                        definition_id=data["_dku_grp_def_id"],
                        edge_dataset="",
                        edge_group=data["_label"],
                        edge_id=data["_dku_grp_id"],
                        source_node_id=source_id,
                        target_node_id=target_id,
                        source_column=graph.nodes[source]["_dku_reserved_user_defined_pk_prop"],
                        target_column=graph.nodes[target]["_dku_reserved_user_defined_pk_prop"],
                        source_node_group=graph.nodes[source]["_label"],
                        target_node_group=graph.nodes[target]["_label"],
                        property_list=[k for k in data if not k.startswith("_")],
                        filters_association="and",
                        filters_stored=[],
                    )
                )

    def __execute_kuzu_query_graph(self, query: str) -> MultiDiGraph:
        result: QueryResult | None = None
        with self.db_handler.get_new_db_instance(readonly=True).get_new_conn() as conn:
            try:
                result = conn.execute(query)
                graph = result.get_as_networkx()
            except Exception as ex:
                logger.info(f"Failed to execute query {query}.")
                raise Exception(str(ex))
            finally:
                if result:
                    try:
                        result.close()
                    except Exception as ex:
                        logger.exception("Failed to close the query result.")
        return graph

    def __execute_kuzu_query_df(self, query: str) -> DataFrame:
        result: QueryResult | None = None
        with self.db_handler.get_new_db_instance(readonly=True).get_new_conn() as conn:
            try:
                result = conn.execute(query)
                df = result.get_as_df()
            except Exception as ex:
                logger.info(f"Failed to execute query {query}.")
                raise Exception(str(ex))
            finally:
                if result:
                    try:
                        result.close()
                    except Exception as ex:
                        logger.exception("Failed to close the query result.")
        return df
