from __future__ import annotations

import logging
from typing import Any, Dict, List, TypedDict

from networkx import MultiDiGraph
from pandas import DataFrame

from graph.connections.db_instance import DbInstance
from graph.connections.query_result import QueryResult
from graph.ingestion.graph_builder import GROUP_DEFINITION_ID_PROP_NAME
from solutions.graph.models import (
    EdgeGroupDefinition,
    EdgeGroupId,
    GraphDBDoesNotExistError,
    NodeGroupDefinition,
    NodeGroupId,
)

logger = logging.getLogger(__name__)


class CypherQueryExecutionException(Exception):
    """
    The goal of this exception is to discriminate between errors occuring because of invalid user queries and actual misbehaviors from VGE.
    """

    pass


class UserQueryResult(TypedDict):
    graph: MultiDiGraph
    table: DataFrame


class UserQueryExecutor:
    def __init__(
        self,
        db_instance: DbInstance,
    ) -> None:
        self.__db_instance = db_instance

    def execute(
        self, query: str, params: Dict[str, Any] | None = None, timeout_seconds: int | None = None
    ) -> UserQueryResult:
        """
        Raises:
            GraphDatabaseDoesNotExistError
            GraphDBWriteInProgressError
            GraphDBReadInProgressError
            CypherQueryExecutionException
        """
        if not self.__db_instance.exists():
            raise GraphDBDoesNotExistError

        with self.__db_instance.get_new_conn(timeout_seconds) as connection:
            result: QueryResult | None = None
            try:
                result = connection.execute(query, params)
                graph: MultiDiGraph = result.get_as_networkx()

                return UserQueryResult(graph=graph, table=result.get_as_df())
            except Exception as ex:
                logger.info(f"Failed to execute query {query}.")
                raise CypherQueryExecutionException(str(ex))
            finally:
                if result:
                    try:
                        result.close()
                    except Exception as ex:
                        logger.exception("Failed to close the query result.")


class EnrichGraphWithConfigurationMetadata:
    def __init__(
        self,
        node_definitions: List[NodeGroupDefinition],
        edge_definitions: List[EdgeGroupDefinition],
        graph: MultiDiGraph,
    ) -> None:
        """
        Enrich a NetworkX graph with configuration metadata from node and edge definitions.
        
        This class modifies the provided graph in-place by adding metadata attributes
        to nodes and edges based on their definitions. The enriched graph organizes
        raw database properties into structured metadata for easier querying.
        
        Args:
            node_definitions: List of node group definitions containing metadata.
            edge_definitions: List of edge group definitions containing metadata.
            graph: NetworkX MultiDiGraph to enrich (modified in-place).
        """
        self.__graph = graph

        self.__load_node_groups_from_graph__(node_definitions)
        self.__load_edge_groups_from_graph__(edge_definitions)

    def get_nodes(self, node_id: NodeGroupId) -> List:
        """
        Retrieve all nodes belonging to a specific node group.
        
        Args:
            node_id: The node group ID to filter by.
            
        Returns:
            List of node dictionaries containing enriched metadata and properties.
        """
        return [self.__graph.nodes[id] for id, attr in self.__graph.nodes(data=True) if attr.get("group") == node_id]

    def get_edges(self, edge_id: EdgeGroupId) -> List:
        """
        Retrieve all edges belonging to a specific edge group.
        
        Args:
            edge_id: The edge group ID to filter by.
            
        Returns:
            List of edge dictionaries containing enriched metadata and properties,
            including the edge key from the MultiDiGraph.
        """
        return [
            {"key": key, **data}
            for _, _, key, data in self.__graph.edges(data=True, keys=True)  # type: ignore
            if data["group"] == edge_id
        ]

    def __load_node_groups_from_graph__(
        self, group_definitions: List[NodeGroupDefinition]
    ) -> None:
        """
        Enrich nodes in the graph with metadata from their definitions.
        
        This method modifies nodes in-place by:
        - Matching each node to its definition via GROUP_DEFINITION_ID_PROP_NAME
        - Extracting the label from the configured label column
        - Clearing existing attributes and replacing with structured metadata
        - Moving all node properties into a 'properties' sub-dictionary
        
        Args:
            group_definitions: List of node group definitions to enrich from.
        """
        for node_instance_id, data in list(self.__graph.nodes(data=True)):
            node_definition_id = data[GROUP_DEFINITION_ID_PROP_NAME]
            node_definition = next(
                (definition for definition in group_definitions if definition["definition_id"] == node_definition_id),
                None,
            )
            assert node_definition

            label_column = node_definition["label_col"]
            node_group = node_definition["node_group"]

            properties_data = dict(data)

            # Remove internal properties
            del properties_data["_id"]
            del properties_data["_label"]

            self.__graph.nodes[node_instance_id].clear()

            self.__graph.nodes[node_instance_id].update({
                "label": properties_data[label_column],
                "id": node_instance_id,
                "name": node_group,
                "group": node_definition["node_id"],
                "properties": properties_data,
            })


    def __load_edge_groups_from_graph__(
        self, group_definitions: List[EdgeGroupDefinition]
    ) -> None:
        """
        Enrich edges in the graph with metadata from their definitions.
        
        This method modifies edges in-place by:
        - Matching each edge to its definition via GROUP_DEFINITION_ID_PROP_NAME
        - Filtering properties to only those specified in the definition
        - Clearing existing attributes and replacing with structured metadata
        - Adding source/target references and moving properties into a sub-dictionary
        
        Args:
            group_definitions: List of edge group definitions to enrich from.
        """
        for source_node_instance_id, target_node_instance_id, key, data in list(self.__graph.edges(data=True, keys=True)):  # type: ignore
            edge_definition_id = data[GROUP_DEFINITION_ID_PROP_NAME]
            edge_definition = next(
                (definition for definition in group_definitions if definition["definition_id"] == edge_definition_id),
                None,
            )
            assert edge_definition

            edge_instance = {k: v for k, v in data.items() if k in edge_definition["property_list"]}

            edge_data = self.__graph[source_node_instance_id][target_node_instance_id][key]
            edge_data.clear()

            edge_data.update({
                "id": source_node_instance_id + target_node_instance_id,
                "group": edge_definition["edge_id"],
                "from": source_node_instance_id,
                "to": target_node_instance_id,
                "properties": edge_instance,
            })
