from __future__ import annotations

import logging

from flask import Blueprint, request

from backends.explorer.connections.db_handler_registry import get_handler_for_snapshot
from backends.explorer.metadata.metadata_manager import ExplorerMetadataManager
from backends.explorer.utils.webapp_config import webapp_config
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

logger = logging.getLogger(__name__)

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


@graph_queries.route("/<snapshot_id>/computeAdjacentNodeGroupsInfo", methods=["POST"])
def get_adjacent_node_groups(snapshot_id: str):
    params = ComputeAdjacentNodeGroupInfoParams(**request.get_json())
    db_handler = get_handler_for_snapshot(snapshot_id, webapp_config.db_explorer_folders)
    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("/<snapshot_id>/getAdjacentNodes", methods=["POST"])
def get_adjacent_nodes_endpoint(snapshot_id: str):
    params = GetAdjacentNodesParams(**request.get_json())
    
    db_handler = get_handler_for_snapshot(snapshot_id, webapp_config.db_explorer_folders)
    graph_manager = ExplorerMetadataManager(webapp_config.db_explorer_folders)
    explorer_metadata = graph_manager.get_snapshot_metadata(snapshot_id, db_handler)
    
    node_definitions = explorer_metadata["node_definitions"]
    edge_definitions = explorer_metadata["edge_definitions"]
    
    with db_handler.get_new_db_instance(readonly=True) as db_instance:
        return return_ok(get_adjacent_nodes(db_instance, params, node_definitions, edge_definitions))


@graph_queries.route("/<snapshot_id>/getTotalCounts", methods=["GET"])
def get_total_counts(snapshot_id: str):
    db_handler = get_handler_for_snapshot(snapshot_id, webapp_config.db_explorer_folders)
    with db_handler.get_new_db_instance(readonly=True) as db_instance:
        return return_ok(compute_total_counts(db_instance))
