import logging
from typing import List

from backends.shared.params import ComputeAdjacentNodeGroupInfoParams, GetAdjacentNodesParams
from backends.shared.response import GetAdjacentNodesSuccess
from graph.connections.db_instance import DbInstance
from graph.exploration.cypher_graph_view import EnrichGraphWithConfigurationMetadata
from graph.exploration.queries.node_adjacency import (
    AdjacentNodeGroupInfo,
    get_adjacent_node_groups_info,
)
from graph.exploration.queries.node_adjacency import (
    get_adjacent_nodes as get_adjacent_nodes_query,
)
from graph.exploration.queries.total_counts import (
    GraphTotalCounts,
    get_edge_counts_by_group,
    get_node_counts_by_group,
)
from graph.exploration.schema.schema_extractor import GraphSchema, extract_schema
from solutions.graph.models import EdgeGroupDefinition, NodeGroupDefinition

logger = logging.getLogger(__name__)


def compute_adjacent_node_groups_info(
    db_instance: DbInstance, params: ComputeAdjacentNodeGroupInfoParams
) -> List[AdjacentNodeGroupInfo]:
    schema: GraphSchema = extract_schema(db_instance)
    node_group_schema = schema["node_groups"].get(params.node_group)
    if not node_group_schema:
        raise Exception(f'Node group "{params.node_group}" does not exist in this graph database.')

    with db_instance.get_new_conn() as connection:
        return get_adjacent_node_groups_info(connection, node_group_schema, params.node_id)


def get_adjacent_nodes(
    db_instance: DbInstance,
    params: GetAdjacentNodesParams,
    node_group_definitions: List[NodeGroupDefinition],
    edge_group_definitions: List[EdgeGroupDefinition],
) -> GetAdjacentNodesSuccess:
    schema: GraphSchema = extract_schema(db_instance)
    source_node_group_schema = schema["node_groups"].get(params.source_node_group)
    if not source_node_group_schema:
        raise ValueError(f'Node group "{params.source_node_group}" does not exist in this graph database.')

    target_node_group_schema = None
    if params.target_node_group:
        target_node_group_schema = schema["node_groups"].get(params.target_node_group)
        if not target_node_group_schema:
            raise ValueError(f'Node group "{params.target_node_group}" does not exist in this graph database.')

    with db_instance.get_new_conn() as connection:
        graph = get_adjacent_nodes_query(connection, source_node_group_schema, params.node_id, target_node_group_schema)
        graph_enricher = EnrichGraphWithConfigurationMetadata(node_group_definitions, edge_group_definitions, graph)

        nodes = []
        for node_def in node_group_definitions:
            nodes.extend(graph_enricher.get_nodes(node_def["node_id"]))

        edges = []
        for edge_def in edge_group_definitions:
            edges.extend(graph_enricher.get_edges(edge_def["edge_id"]))

        return {"nodes": nodes, "edges": edges}


def compute_total_counts(db_instance: DbInstance) -> GraphTotalCounts:
    if not db_instance.exists():
        logger.info("Graph database doesn't exist, ignore node count request")
        node_counts, edge_counts = {}, {}
    else:
        with db_instance.get_new_conn() as connection:
            node_counts = get_node_counts_by_group(connection)
            edge_counts = get_edge_counts_by_group(connection)

    return {"nodes": node_counts, "edges": edge_counts}
