from __future__ import annotations

import logging
import time
from typing import Any, Dict, Optional

from backends.shared.params import RunCypherParams, RunLlmCypherParams
from backends.shared.response import RunCypherFailure, RunCypherSuccess
from graph.connections.db_handler import DbHandler
from graph.exploration.cypher_graph_view import (
    CypherQueryExecutionException,
    EnrichGraphWithConfigurationMetadata,
    UserQueryExecutor,
)
from solutions.graph.models import (
    GraphDBDoesNotExistError,
    GraphDBWriteInProgressError,
    to_edge_group_definitions,
    to_node_group_definitions,
)
from stores.graph_metadata_store import GraphMetadataStore

logger = logging.getLogger(__name__)


def run_cypher(
    store: GraphMetadataStore,
    db_handler: DbHandler,
    params: RunCypherParams | RunLlmCypherParams,
    timeout_seconds: int | None = None,
) -> RunCypherSuccess | RunCypherFailure:
    """
    Raises:
        GraphDoesNotExistError
        GraphMetadataStoreError
        GraphDBWriteInProgressError
        GraphDBDoesNotExistError
    """
    perf_t0 = time.perf_counter()
    query_params: Optional[Dict[str, Any]] = None
    if isinstance(params, RunCypherParams):
        query_params = params.params
    graph_id = params.graph_id
    graph_metadata = 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:
        user_query_executor = UserQueryExecutor(db_instance)

        try:
            perf_t1 = time.perf_counter()
            user_query_result = user_query_executor.execute(params.query, query_params, timeout_seconds)
            graph_enricher = EnrichGraphWithConfigurationMetadata(
                node_group_definitions, edge_group_definitions, user_query_result["graph"], user_query_result["table"]
            )
            perf_t2 = time.perf_counter()
        except CypherQueryExecutionException as ex:
            logger.info(f"User query raised an error on graph {graph_id}, '{params.query}'.", exc_info=True)
            return {"success": False, "errorCode": "INVALID_CYPHER_ERROR", "error": str(ex)}
        except GraphDBWriteInProgressError as ex:
            raise ex
        except GraphDBDoesNotExistError as ex:
            raise ex
        except RuntimeError as ex:
            logger.exception(f"Runtime error on graph {graph_id}, '{params.query}'.", exc_info=True)
            return {
                "success": False,
                "errorCode": "BROKEN_DB_STATE_ERROR",
                "error": "The underlying graph database seems to be in a broken state. You may consider rebuilding it.",
            }
        except Exception as ex:
            logger.exception(f"An unexpected error occured on graph {graph_id}, '{params.query}'.", exc_info=True)
            return {"success": False, "errorCode": "UNEXPECTED_ERROR", "error": "An unexpected error occured."}

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

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

        table = graph_enricher.get_table()

        perf_t3 = time.perf_counter()

        logger.debug(f"Cypher query timing: {(perf_t3 - perf_t0):.2f}s total")
        logger.debug(f".........query prep: {(perf_t1 - perf_t0):.2f}s")
        logger.debug(f"....query execution: {(perf_t2 - perf_t1):.2f}s")
        logger.debug(f"(configured timeout: {timeout_seconds}s)")
        logger.debug(f"....post-processing: {(perf_t3 - perf_t2):.2f}s")

        return {
            "success": True,
            "nodes": nodes,
            "edges": edges,
            "table": {"columns": table["columns"], "rows": table["df"].to_dict("records")} if table else None,
        }
