from __future__ import annotations

import logging

from flask import Blueprint, request

from backends.editor.utils.api import use_db_handler, use_graph_metadata_store
from backends.shared.params import ComputeAdjacentNodeGroupInfoParams, GetAdjacentNodesParams
from backends.shared.queries import compute_adjacent_node_groups_info, compute_total_counts, get_adjacent_nodes
from backends.utils import return_ok
from graph.connections.db_handler import DbHandler
from solutions.graph.models import to_edge_group_definitions, to_node_group_definitions
from stores.graph_metadata_store import GraphMetadataStore

logger = logging.getLogger(__name__)

graph_queries = Blueprint("graph_queries", __name__, url_prefix="/queries")


@graph_queries.route("/<graph_id>/computeAdjacentNodeGroupsInfo", methods=["POST"])
@use_db_handler(id_getter=lambda: request.view_args.get("graph_id") if request.view_args else None)
def get_adjacent_node_groups(graph_id: str, db_handler: DbHandler):
    params = ComputeAdjacentNodeGroupInfoParams(**request.get_json())

    with db_handler.get_new_db_instance(readonly=True) as db_instance:
        return return_ok(compute_adjacent_node_groups_info(db_instance, params))


@graph_queries.route("/<graph_id>/getAdjacentNodes", methods=["POST"])
@use_graph_metadata_store
@use_db_handler(id_getter=lambda: request.view_args.get("graph_id") if request.view_args else None)
def get_adjacent_nodes_endpoint(graph_id: str, graph_store: GraphMetadataStore, db_handler: DbHandler):
    params = GetAdjacentNodesParams(**request.get_json())
    graph_metadata = graph_store.get(graph_id)

    node_group_definitions = [
        definition
        for node_group_metadata in graph_metadata["nodes"].values()
        for definition in to_node_group_definitions(node_group_metadata)
    ]
    edge_group_definitions = [
        definition
        for edge_group_metadata in graph_metadata["edges"].values()
        for definition in to_edge_group_definitions(edge_group_metadata, graph_metadata)
    ]
    with db_handler.get_new_db_instance(readonly=True) as db_instance:
        return return_ok(get_adjacent_nodes(db_instance, params, node_group_definitions, edge_group_definitions))


@graph_queries.route("/<graph_id>/getTotalCounts", methods=["GET"])
@use_db_handler(id_getter=lambda: request.view_args.get("graph_id") if request.view_args else None)
def get_total_counts(graph_id: str, db_handler: DbHandler):
    with db_handler.get_new_db_instance(readonly=True) as db_instance:
        return return_ok(compute_total_counts(db_instance))
