import logging

from flask import Blueprint, g

from backends.explorer.queries.explorer_data import run_llm_query
from backends.explorer.utils.connections import get_webapp_db_handler
from backends.explorer.utils.webapp_config import webapp_config
from backends.shared.params import RunLlmCypherParams
from backends.utils import return_ok, set_auth_identifier, validate_params
from graph.exploration.llm import AIQueryExecutionException
from graph.exploration.llm.logging_dataset import logging_dataset
from solutions.graph.models import FeatureUnavailableError, PresetNotSetException

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"])
@validate_params(RunLlmCypherParams)
def run_llm(params: RunLlmCypherParams):
    """
    Generates and executes a Cypher query from a natural language question using an LLM.

    The generated query is executed against the database, and results are returned
    to the user. If initialized, the query and answer are logged to a DSS dataset.

    Args:
        params (RunLlmCypherParams): Validated parameters including graph ID and query.

    Returns:
        dict: OK response with execution results or error details.
    """
    auth_identifier = g.get("auth_identifier")

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

    try:
        db_handler, conn_info_registry = get_webapp_db_handler(params.graph_id)

        data = run_llm_query(
            params,
            llm,
            db_handler,
            conn_info_registry,
            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 PresetNotSetException as ex:
        raise ex
    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.",
            }
        )

