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, TableData
from graph.ingestion.dataiku.data_types import GraphDataikuDataType
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: TableData


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()
                table: TableData = result.get_as_table()

                return UserQueryResult(graph=graph, table=table)
            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 TableColumn(TypedDict):
    name: str
    type: GraphDataikuDataType


class Table(TypedDict):
    columns: list[TableColumn]
    df: DataFrame


class EnrichGraphWithConfigurationMetadata:
    def __init__(
        self,
        node_definitions: List[NodeGroupDefinition],
        edge_definitions: List[EdgeGroupDefinition],
        graph: MultiDiGraph,
        table: TableData | None = None,
    ) -> 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: MultiDiGraph representing the graph structure.
            table: TableData containing the table data, if available.
        """
        self.__node_definitions_by_id = {definition["definition_id"]: definition for definition in node_definitions}
        self.__edge_definitions_by_id = {definition["definition_id"]: definition for definition in edge_definitions}

        self.__table = table
        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 get_table(self) -> Table | None:
        if not self.__table:
            return None

        schema = self.__table["schema"]
        return {
            "df": self.__format_dataframe__(self.__table),
            "columns": [{"name": col, "type": schema[col]} for col in schema.keys()],
        }

    def __format_dataframe__(self, table: TableData) -> DataFrame:
        schema = table["schema"]
        df = table["df"].copy()

        for schema_column, column_type in schema.items():
            if schema_column not in df.columns:
                continue

            if column_type == "node":
                df[schema_column] = df[schema_column].apply(self.__format_node_table_cell__)
            elif column_type == "edge":
                df[schema_column] = df[schema_column].apply(self.__format_edge_table_cell__)

        return df.fillna("")

    def __get_node_definition__(self, node_definition_id: str) -> NodeGroupDefinition:
        node_definition = self.__node_definitions_by_id.get(node_definition_id)
        if not node_definition:
            raise ValueError(f"Node definition with id {node_definition_id} not found.")
        return node_definition

    def __get_node_table_properties__(
        self, node_data: dict[str, Any], node_definition: NodeGroupDefinition
    ) -> dict[str, Any]:
        columns_to_keep = set(node_definition["property_list"] + [node_definition["label_col"]])
        return {key: value for key, value in node_data.items() if key in columns_to_keep}

    def __format_node_table_cell__(self, cell: Any) -> Any:
        if not isinstance(cell, dict):
            return cell

        node_definition_id = cell.get(GROUP_DEFINITION_ID_PROP_NAME)
        if not node_definition_id:
            return cell

        node_definition = self.__get_node_definition__(node_definition_id)
        pk_col = node_definition["primary_col"]
        label_column = node_definition["label_col"]
        node_group = node_definition["node_group"]
        node_instance_id = cell.get(pk_col)
        properties_data = self.__get_node_table_properties__(cell, node_definition)

        return {
            "label": properties_data.get(label_column),
            "id": node_instance_id,
            "name": node_group,
            "group": node_definition["node_id"],
            "properties": properties_data,
        }

    def __get_edge_definition__(self, edge_definition_id: str) -> EdgeGroupDefinition:
        edge_definition = self.__edge_definitions_by_id.get(edge_definition_id)
        if not edge_definition:
            raise ValueError(f"Edge definition with id {edge_definition_id} not found.")
        return edge_definition

    def __get_edge_table_properties__(
        self, edge_data: dict[str, Any], edge_definition: EdgeGroupDefinition
    ) -> dict[str, Any]:
        columns_to_keep = set(edge_definition["property_list"])
        return {key: value for key, value in edge_data.items() if key in columns_to_keep}

    def __format_edge_table_cell__(self, cell: Any) -> Any:
        if not isinstance(cell, dict):
            return cell

        edge_definition_id = cell.get(GROUP_DEFINITION_ID_PROP_NAME)
        if not edge_definition_id:
            return cell

        edge_definition = self.__get_edge_definition__(edge_definition_id)
        source_node_instance_id = cell.get("_src")
        target_node_instance_id = cell.get("_dst")
        properties_data = self.__get_edge_table_properties__(cell, edge_definition)

        return {
            "id": f"{source_node_instance_id}~{target_node_instance_id}",
            "name": edge_definition["edge_group"],
            "group": edge_definition["edge_id"],
            "from": source_node_instance_id,
            "to": target_node_instance_id,
            "properties": properties_data,
        }

    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 = self.__get_node_definition__(node_definition_id)

            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 = self.__get_edge_definition__(edge_definition_id)

            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": f"{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,
                }
            )
