from __future__ import annotations

import logging
from typing import Any

from networkx import MultiDiGraph
from pandas import DataFrame

from backends.explorer.metadata.metadata_builder import Colorizer, MetadataBuilder
from graph.connections.db_handler import DbHandler
from graph.connections.query_result import QueryResult
from graph.ingestion.graph_builder import GROUP_DEFINITION_ID_PROP_NAME
from solutions.graph.models import EdgeGroupDefinition, ExplorerMetadata, NodeGroupDefinition

logger = logging.getLogger(__name__)


class Neo4jMetadataBuilder(MetadataBuilder):
    def __init__(self, db_handler: DbHandler, snapshot_id: str):
        super().__init__(db_handler, snapshot_id)

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

        # Get all node labels
        node_labels = self.__get_node_labels()
        for label in node_labels:
            self._process_node_label(label, node_definitions)

        # Get all relationship types
        rel_types = self.__get_relationship_types()
        for rel_type in rel_types:
            self._process_relationship_type(rel_type, edge_definitions)

        return ExplorerMetadata(
            snapshot_id=self.snapshot_id,
            db_type="neo4j",
            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 __get_node_labels(self) -> list[str]:
        """Get all node labels from the Neo4j database."""
        query = "CALL db.labels()"
        df = self.__execute_query_df(query)
        return df["label"].tolist() if not df.empty else []

    def __get_relationship_types(self) -> list[str]:
        """Get all relationship types from the Neo4j database."""
        query = "CALL db.relationshipTypes()"
        df = self.__execute_query_df(query)
        return df["relationshipType"].tolist() if not df.empty else []
    
    def _get_default_node_label(self, properties: dict[str, Any]) -> str:
        """
        Determine the default label property for a node based on a prioritized list of property keys.
        
        Args:
            properties: A dictionary of node properties.
            
        Returns:
            The name of the property to use as the default label.
        """

        # To get the default node caption we use kinda the same way than neo4j brower
        # https://github.com/neo4j/neo4j-browser/blob/ff8ed858f50b095d00c0fba725487f192900b964/src/neo4j-arc/graph-visualization/models/GraphStyle.ts#L283
        priority_keys = ["name", "title", "label", "code", "id"]

        keys_map = {k.lower(): k for k in properties.keys() if not k.startswith("_")}

        for candidate in priority_keys:
            if candidate in keys_map:
                return str(keys_map[candidate])

        return "_id"

    def _process_node_label(self, label: str, node_definitions: list[NodeGroupDefinition]) -> None:
        """Process a node label to extract node definitions."""
        # Query to get a sample node of this label
        query = f"MATCH (o:`{label}`) RETURN o LIMIT 1"
        graph = self.__execute_query_graph(query)
        for node, data in graph.nodes(data=True):
            # Take the label if GROUP_DEFINITION_ID_PROP_NAME because graph was built outside of VG
            def_id = data.get(GROUP_DEFINITION_ID_PROP_NAME, label)
            label_col = self._get_default_node_label(data)
            properties = [k for k in data if not k.startswith("_")]
            node_definitions.append(
                NodeGroupDefinition(
                    definition_id=def_id,
                    node_group=label,
                    node_id=node,
                    source_dataset="",
                    primary_col="id",
                    label_col=label_col,
                    property_list=properties,
                    filters_association="and",
                    filters_stored=[],
                )
            )

    def _process_relationship_type(self, rel_type: str, edge_definitions: list[EdgeGroupDefinition]) -> None:
        """Process a relationship type to extract edge definitions."""
        # Query to get a sample relationship of this type
        query = f"MATCH (o)-[e:`{rel_type}`]->(d) RETURN o, e, d LIMIT 1"
        graph = self.__execute_query_graph(query)
        for source, target, data in graph.edges(data=True):
            source_column = "id" if "id" in data else "_src"
            target_column = "id" if "id" in data else "_dst"
            # Take the rel type if GROUP_DEFINITION_ID_PROP_NAME because graph was built outside of VG
            def_id = data.get(GROUP_DEFINITION_ID_PROP_NAME, rel_type)
            edge_definitions.append(
                EdgeGroupDefinition(
                    definition_id=def_id,
                    edge_dataset="",
                    edge_group=rel_type,
                    edge_id=rel_type,
                    source_node_id=source,
                    target_node_id=target,
                    source_column=source_column,
                    target_column=target_column,
                    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 _process_col(
        self, col: dict, node_definitions: list[NodeGroupDefinition], edge_definitions: list[EdgeGroupDefinition]
    ) -> None:
        name = col["name"]

        if col["type"] == "NODE":
            query = f"MATCH (o:`{name}`) RETURN o LIMIT 1"
            graph = self.__execute_query_graph(query)
            for node, data in graph.nodes(data=True):
                label_col = self._get_default_node_label(data)
                properties = [k for k in data if not k.startswith("_")]

                node_definitions.append(
                    NodeGroupDefinition(
                        definition_id=data["_label"],
                        node_group=data["_label"],
                        node_id=node,
                        source_dataset="",
                        primary_col="id",
                        label_col=label_col,
                        property_list=properties,
                        filters_association="and",
                        filters_stored=[],
                    )
                )
        else:
            query = f"MATCH (o)-[e:`{name}`]->(d) RETURN o, e, d LIMIT 1"
            graph = self.__execute_query_graph(query)
            for source, target, data in graph.edges(data=True):
                edge_definitions.append(
                    EdgeGroupDefinition(
                        definition_id=data["_label"],
                        edge_dataset="",
                        edge_group=data["_label"],
                        edge_id=data["_dku_grp_id"],
                        source_node_id=source,
                        target_node_id=target,
                        source_column=graph.nodes[source]["_label"],
                        target_column=graph.nodes[target]["_label"],
                        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_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:
            result = conn.execute(query)
            graph = result.get_as_networkx()
        return graph

    def __execute_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:
            result = conn.execute(query)
            df = result.get_as_df()
        return df
