from __future__ import annotations

from dataikuapi.dss.llm import DSSLLM

from backends.explorer.discovery.info_registry import ConnectionInfoRegistry
from backends.explorer.metadata.metadata_manager import ExplorerMetadataManager
from backends.explorer.queries.params import ExplorerRunCypherParams
from backends.shared.params import RunLlmCypherParams
from backends.shared.response import RunCypherFailure, RunCypherSuccess
from graph.connections.db_handler import DbHandler
from graph.exploration.cypher_graph_view import EnrichGraphWithConfigurationMetadata, UserQueryExecutor
from graph.exploration.llm import LLMCypher
from solutions.graph.models import ExplorerMetadata


def run_cypher_query(
    params: ExplorerRunCypherParams,
    db_handler: DbHandler,
    conn_info_registry: ConnectionInfoRegistry,
    timeout_seconds: int | None = None
) -> RunCypherSuccess | RunCypherFailure:
    """
    Executes a Cypher query against a graph database saved configuration and returns the results.
    Args:
        params (ExplorerRunCypherParams): Parameters for running the Cypher query, including the query string and saved configuration ID.
        db_folder (Folder): The folder containing the database files.
    Returns:
        RunCypherSuccess | RunCypherFailure: A dictionary containing:
            - "success" (bool): Indicates if the query was executed successfully.
            - "nodes" (list): List of nodes returned by the query.
            - "edges" (list): List of edges returned by the query.
            - "table" (dict): Tabular representation of the query result with:
                - "columns" (list): List of column definitions (each with a "name" key).
                - "rows" (list): List of row dictionaries representing the result set.
    Raises:
        Any exceptions raised by the underlying graph manager, metadata builder, or Cypher execution.
    """
    graph_manager = ExplorerMetadataManager(conn_info_registry)

    query_params = params.params

    explorer_metadata = graph_manager.get_snapshot_metadata(params.snapshot_id, db_handler)

    with db_handler.get_new_db_instance(readonly=True) as db_instance:
        user_query_executor = UserQueryExecutor(db_instance)
        user_query_result = user_query_executor.execute(params.query, query_params, timeout_seconds)
        graph_enricher = EnrichGraphWithConfigurationMetadata(
            explorer_metadata["node_definitions"], explorer_metadata["edge_definitions"], user_query_result["graph"]
        )

        # Get unique node IDs and edge IDs in node definitions
        node_ids = set([d["node_id"] for d in explorer_metadata["node_definitions"]])
        edge_ids = set([d["edge_id"] for d in explorer_metadata["edge_definitions"]])

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

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

        df = user_query_result["table"].fillna("")

        return {
            "success": True,
            "nodes": nodes,
            "edges": edges,
            "table": {"columns": [{"name": col} for col in df.columns], "rows": df.to_dict("records")},
        }


def get_snapshot_metadata(snapshot_id: str, db_handler: DbHandler, conn_info_registry: ConnectionInfoRegistry) -> ExplorerMetadata:
    """
    Retrieve the metadata for a specific saved configuration from the database folder.

    Args:
        snapshot_id (str): The unique identifier of the saved configuration whose metadata is to be retrieved.
        db_folder (Folder): The folder object representing the database location.

    Returns:
        ExplorerMetadata: The metadata associated with the specified saved configuration.

    Raises:
        Any exceptions raised by ExplorerMetadataManager.get_snapshot_metadata.
    """
    graph_manager = ExplorerMetadataManager(conn_info_registry)
    return graph_manager.get_snapshot_metadata(snapshot_id, db_handler)


def get_snapshots_metadata_from_folders(conn_info_registry: ConnectionInfoRegistry) -> list[ExplorerMetadata]:
    """
    Retrieves metadata for all graph saved configurations stored in the specified folder.

    Args:
        db_folder (Folder): The folder containing the graph saved configuration data.

    Returns:
        List[ExplorerMetadata]: A list of metadata objects for each saved configuration found in the folder.
    """
    graph_manager = ExplorerMetadataManager(conn_info_registry)
    return graph_manager.get_snapshots_metadata()


def run_llm_query(
    params: RunLlmCypherParams,
    llm: DSSLLM,
    db_handler: DbHandler,
    conn_info_registry: ConnectionInfoRegistry,
    timeout_seconds: int | None = None
) -> dict:
    """
    Executes a query on a graph database using a Large Language Model (LLM) to generate and run Cypher queries.

    Args:
        params (RunLlmCypherParams): Parameters for the LLM Cypher query, including the graph ID and the user query.
        db_folder (Folder): The folder containing the database metadata and files.
        llm (DSSLLM): The LLM instance used to generate Cypher queries from natural language.
        timeout_seconds: Seconds before query times out.

    Returns:
        dict: The result of executing the generated Cypher query, typically containing query results or error information.
    """
    graph_manager = ExplorerMetadataManager(conn_info_registry)

    explorer_metadata = graph_manager.get_snapshot_metadata(params.graph_id, db_handler)

    with db_handler.get_new_db_instance(readonly=True) as db_instance:
        llm_cypher = LLMCypher(
            db_instance, llm, explorer_metadata["node_definitions"], explorer_metadata["edge_definitions"]
        )
        return llm_cypher.execute_llm_with_retry(params.query, timeout_seconds=timeout_seconds)
