from __future__ import annotations

from typing import List

from dataiku import Folder
from dataikuapi.dss.llm import DSSLLM

from solutions.graph.explorer_builder import ExplorerMetadataManager
from solutions.graph.models import ExplorerMetadata
from solutions.graph.queries.params import ExplorerRunCypherParams, RunLlmCypherParams
from solutions.graph.queries.response import RunCypherFailure, RunCypherSuccess
from solutions.graph.views.cypher_graph_view import CypherGraphView
from solutions.llm import LLMCypher


def run_cypher_query(
    params: ExplorerRunCypherParams, db_folders: List[Folder], 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(db_folders)

    query_params = params.params

    explorer_metadata = graph_manager.get_snapshot_metadata(params.snapshot_id)

    db_instance = graph_manager.db_instance_factory.get_db_instance_from_snapshot_id(params.snapshot_id)
    graph_view = CypherGraphView(
        db_instance, explorer_metadata["node_definitions"], explorer_metadata["edge_definitions"]
    )
    graph_view.execute(params.query, query_params, timeout_seconds)

    # 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_view.get_nodes(node_id))

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

    df = graph_view.get_as_df().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_folders: List[Folder]) -> 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(db_folders)
    return graph_manager.get_snapshot_metadata(snapshot_id)


def get_snapshots_metadata_from_folders(db_folders: List[Folder]) -> 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(db_folders)
    return graph_manager.get_snapshots_metadata()


def run_llm_query(
    params: RunLlmCypherParams, db_folders: List[Folder], llm: DSSLLM, 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(db_folders)

    explorer_metadata = graph_manager.get_snapshot_metadata(params.graph_id)

    db_instance = graph_manager.db_instance_factory.get_db_instance_from_snapshot_id(params.graph_id)

    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)
