import logging

from flask import Blueprint, g, request

from backends.editor.utils.api import use_db_handler, use_graph_metadata_store
from backends.editor.utils.webapp_config import webapp_config
from backends.shared.params import RunLlmCypherParams
from backends.utils import return_ok, set_auth_identifier
from graph.connections.db_handler import DbHandler
from graph.exploration.llm import AIQueryExecutionException, LLMCypher
from graph.exploration.llm.logging_dataset import logging_dataset
from solutions.graph.models import (
    FeatureUnavailableError,
    GraphDBDoesNotExistError,
    GraphDoesNotExistError,
    to_edge_group_definitions,
    to_node_group_definitions,
)
from stores.graph_metadata_store import GraphMetadataStore

logger = logging.getLogger(__name__)
llm = Blueprint("llm", __name__, url_prefix="/llm")


@llm.before_request
def llm_before_request():
    set_auth_identifier()


@llm.route("/generate_cypher", methods=["POST"])
@use_graph_metadata_store
@use_db_handler(id_getter=lambda: request.get_json().get("graph_id"))
def run_llm(graph_store: GraphMetadataStore, db_handler: DbHandler):
    params = RunLlmCypherParams(**request.get_json())
    auth_identifier = g.get("auth_identifier")

    llm = webapp_config.get_llm()
    if not llm:
        raise FeatureUnavailableError("LLM must be set in webapp config")

    graph_id = params.graph_id
    if not graph_store.exists(graph_id):
        raise GraphDoesNotExistError(graph_id)

    graph_metadata = graph_store.get(graph_id)
    nodes_definitions = [
        definition
        for node_group_metadata in graph_metadata["nodes"].values()
        for definition in to_node_group_definitions(node_group_metadata)
    ]
    edges_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:
        if not db_instance.exists():
            raise GraphDBDoesNotExistError

        llm_cypher = LLMCypher(
            db_instance,
            llm,
            nodes_definitions,
            edges_definitions,
        )

        try:
            data = llm_cypher.execute_llm_with_retry(
                params.query, timeout_seconds=webapp_config.db_query_timeout_seconds
            )
            if logging_dataset.is_initialized and data["success"]:
                log_uuid = logging_dataset.log_query(
                    user=auth_identifier,
                    title=data["title"],
                    llm_name=llm.llm_id,
                    question=params.query,
                    answer=data["cypher_query"],
                    graph_id=params.graph_id,
                )
                data["log_uuid"] = log_uuid
            else:
                data["log_uuid"] = None
            return return_ok(data=data)

        except AIQueryExecutionException as ex:
            logger.debug(f"AI query raised an error {params.graph_id}, '{ex.query}'.", exc_info=True)
            return return_ok(
                data={
                    "success": False,
                    "errorCode": "AI_QUERY_CYPHER_ERROR",
                    "error": str(ex),
                    "title": ex.title,
                    "cypher_query": ex.query,
                }
            )
        except Exception as ex:
            logger.exception(
                f"An unexpected error occured on AI Query {params.graph_id}, '{params.query}'.", exc_info=True
            )
            return return_ok(
                data={
                    "success": False,
                    "errorCode": "UNEXPECTED_ERROR",
                    "error": "An unexpected error occured.",
                }
            )
