from __future__ import annotations

import json
import logging
import traceback
from typing import List

from dataikuapi.dss.llm import DSSLLM

from graph.connections.db_instance import DbInstance
from graph.connections.kuzu.db_instance import KuzuDbInstance
from graph.connections.neo4j.db_instance import Neo4jDbInstance
from graph.exploration.cypher_graph_view import (
    CypherQueryExecutionException,
    EnrichGraphWithConfigurationMetadata,
    UserQueryExecutor,
    UserQueryResult,
)
from graph.exploration.llm.db_strategy_factory import create_llm_strategy
from graph.exploration.llm.prompts import (
    get_user_prompt,
    kuzu_get_retry_system_prompt,
    kuzu_get_system_prompt,
    neo4j_get_retry_system_prompt,
    neo4j_get_system_prompt,
)
from solutions.graph.models import EdgeGroupDefinition, NodeGroupDefinition

logger = logging.getLogger(__name__)

LLM_RETRY = 3


class AIQueryExecutionException(Exception):
    """
    Exception raised when an error occurs during the execution of an AI-generated Cypher query.

    Attributes:
        query (str): The Cypher query that caused the exception.
        cypher_error (str): The error message returned by the Cypher engine (first line only).

    Args:
        query (str): The generated Cypher query that failed.
        cypher_error (str): The error message from the Cypher engine.

    Example:
        raise AIQueryExecutionException(query, cypher_error)
    """

    def __init__(self, query: str, title: str, cypher_error: str):
        self.query = query
        self.title = title
        self.cypher_error = cypher_error

        message = self.cypher_error

        super().__init__(message)


class LLMCypher:
    def __init__(
        self,
        db_instance: DbInstance,
        llm: DSSLLM,
        node_definitions: List[NodeGroupDefinition],
        edge_definitions: List[EdgeGroupDefinition],
    ):
        self.__db_instance = db_instance
        self.__llm = llm

        self.__node_definitions = node_definitions
        self.__edge_definitions = edge_definitions
        self.__user_query_executor = UserQueryExecutor(db_instance)

        self.__graph_schema: str | None = None

        self.__db_strategy = create_llm_strategy(db_instance)

    def execute_llm_with_retry(
        self,
        question: str,
        cypher_query: str | None = None,
        title: str | None = None,
        retry: int = LLM_RETRY,
        timeout_seconds: int | None = None,
        last_cypher_exception: AIQueryExecutionException | None = None,
    ) -> dict:
        """
        Executes a Cypher query, with retries on failure.

        If a query is not provided, it generates one from the natural language question.

        Args:
            question (str): The natural language question.
            cypher_query (str | None, optional): An optional Cypher query to execute.
            title (str | None, optional): An optional title for the query.
            retry (int, optional): Number of retries allowed. Defaults to LLM_RETRY.
            timeout_seconds (int, optional): Query timeout in seconds. Defaults to None.
            last_cypher_exception (AIQueryExecutionException optional): The last exception encountered during Cypher execution, used for error reporting. Defaults to None.

        Returns:
            dict: The formatted result, including nodes, edges, a table, the query, and title.

        Raises:
            Exception: If all retries are exhausted.
        """
        if retry < 0:
            if isinstance(last_cypher_exception, AIQueryExecutionException):
                raise last_cypher_exception
            else:
                raise Exception(f"Error with query after multiple retries.\n{traceback.format_exc()}")

        current_title = title
        current_cypher_query = cypher_query

        if not current_cypher_query:
            try:
                llm_response = self.generate_cypher(question)
                current_cypher_query = llm_response["cypher_query"]
                current_title = llm_response["title"]
            except (json.JSONDecodeError, KeyError) as e:
                logger.error(f"Failed to parse LLM response: {e}. Retrying...")
                # Retry without a "previous query" since the response was malformed
                return self.execute_llm_with_retry(question=question, retry=retry - 1, timeout_seconds=timeout_seconds)

        if not current_title:
            current_title = "User-provided query"

        try:
            logger.debug(f"Execute cypher query: {current_cypher_query}")
            try:
                user_query_result = self.__user_query_executor.execute(
                    current_cypher_query, timeout_seconds=timeout_seconds
                )
            except CypherQueryExecutionException as e:
                raise AIQueryExecutionException(query=current_cypher_query, title=current_title, cypher_error=str(e))

            return self.__format_graph_view_result(current_cypher_query, current_title, user_query_result)

        except Exception as e:
            error_trace = traceback.format_exc()
            logger.debug(f"Cypher query failed, retrying. Retry: {retry}.\n{error_trace}")

            try:
                llm_response = self.generate_cypher(question, current_cypher_query, error_trace)
                new_cypher_query = llm_response["cypher_query"]
                new_title = llm_response["title"]
                # Recursively call with the new query and title
                return self.execute_llm_with_retry(
                    question=question,
                    cypher_query=new_cypher_query,
                    title=new_title,
                    retry=retry - 1,
                    timeout_seconds=timeout_seconds,
                    last_cypher_exception=e if isinstance(e, AIQueryExecutionException) else None,
                )
            except (json.JSONDecodeError, KeyError) as e:
                logger.error(f"Failed to parse LLM response on retry: {e}. Retrying...")
                return self.execute_llm_with_retry(question=question, retry=retry - 1, timeout_seconds=timeout_seconds)

    def generate_cypher(
        self, llm_request: str, previous_query: str | None = None, error_message: str | None = None
    ) -> dict[str, str]:
        """
        Generates a Cypher query and a title using an LLM.

        Args:
            llm_request (str): The user's request.
            previous_query (str | None, optional): The previous failed query for context.
            error_message (str | None, optional): The error message from the failed query.

        Returns:
            dict: A dictionary with "title" and "cypher_query" keys.

        Raises:
            json.JSONDecodeError: If the LLM response is not valid JSON.
            KeyError: If the JSON response is missing required keys.
        """
        if not self.__graph_schema:
            self.__graph_schema = self.__db_strategy.get_graph_schema()

        if previous_query and error_message:
            system_prompt = self._get_retry_system_prompt(
                self.__graph_schema, llm_request, previous_query, error_message
            )
        else:
            system_prompt = self._get_system_prompt()

        user_prompt = get_user_prompt(self.__graph_schema, llm_request)
        completion = self.__llm.new_completion()

        logger.debug(f"Generate cypher from system prompt: {system_prompt}")
        completion.with_message(system_prompt, role="system")

        logger.debug(f"Generate cypher from user prompt: {user_prompt}")
        completion.with_message(user_prompt, role="user")

        resp = completion.execute()

        logger.debug(f"LLM raw response: {resp.text}")
        parsed_response: dict[str, str] = json.loads(resp.text)  # type: ignore

        if "title" not in parsed_response or "cypher_query" not in parsed_response:
            raise KeyError("LLM response is missing 'title' or 'cypher_query' key.")

        return parsed_response

    def _get_system_prompt(self) -> str:
        """
        Get the system prompt appropriate for the database type.

        Returns:
            str: The system prompt for the LLM.
        """

        if isinstance(self.__db_instance, Neo4jDbInstance):
            return neo4j_get_system_prompt()
        elif isinstance(self.__db_instance, KuzuDbInstance):
            return kuzu_get_system_prompt()
        else:
            raise ValueError(f"Unsupported database type: {type(self.__db_instance).__name__}")

    def _get_retry_system_prompt(self, schema: str, question: str, query: str, error_message: str) -> str:
        """
        Get the retry system prompt appropriate for the database type.

        Args:
            schema (str): The graph schema.
            question (str): The original user question.
            query (str): The failed query.
            error_message (str): The error message.

        Returns:
            str: The retry system prompt for the LLM.
        """

        if isinstance(self.__db_instance, Neo4jDbInstance):
            return neo4j_get_retry_system_prompt(schema, question, query, error_message)
        elif isinstance(self.__db_instance, KuzuDbInstance):
            return kuzu_get_retry_system_prompt(schema, question, query, error_message)
        else:
            raise ValueError(f"Unsupported database type: {type(self.__db_instance).__name__}")

    def __format_graph_view_result(self, cypher_query: str, title: str, user_query_result: UserQueryResult) -> dict:
        """
        Formats the graph view result into a structured dictionary.

        Args:
            cypher_query (str): The Cypher query used.
            title (str): The descriptive title for the query.

        Returns:
            dict: A dictionary containing the query results and metadata.
        """

        # Get unique node IDs and edge IDs in node definitionsAdd commentMore actions
        node_ids = set([d["node_id"] for d in self.__node_definitions])
        edge_ids = set([d["edge_id"] for d in self.__edge_definitions])

        enriched_graph = EnrichGraphWithConfigurationMetadata(
            self.__node_definitions, self.__edge_definitions, user_query_result["graph"], user_query_result["table"]
        )

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

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

        table = enriched_graph.get_table()

        return {
            "success": True,
            "title": title,
            "nodes": nodes,
            "edges": edges,
            "table": {"columns": table["columns"], "rows": table["df"].to_dict("records")} if table else None,
            "cypher_query": cypher_query,
        }
