import logging

from flask import Blueprint, g, request

from explorer.backend.utils.webapp_config import webapp_config
from solutions.backend.utils import return_ok, set_auth_identifier
from solutions.graph.models import FeatureUnavailableError
from solutions.graph.queries.explorer_data import run_llm_query
from solutions.graph.queries.params import RunLlmCypherParams
from solutions.llm import AIQueryExecutionException
from solutions.llm.logging_dataset import logging_dataset

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"])
def run_llm():
    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")

    try:
        data = run_llm_query(
            params, webapp_config.db_explorer_folders, llm, 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.info(
            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.",
        })
