from __future__ import annotations

import logging
from abc import ABC, abstractmethod

from pandas import DataFrame

from graph.connections.db_instance import DbInstance
from graph.connections.query_result import QueryResult
from graph.exploration.cypher_graph_view import CypherQueryExecutionException
from graph.exploration.schema.schema_extractor import GraphSchema, extract_schema

logger = logging.getLogger(__name__)


class LLMCypherDbStrategy(ABC):
    """Abstract base class for DB-specific Cypher operations"""

    def __init__(self, db_instance: DbInstance):
        self.db_instance = db_instance

    def get_graph_schema(self) -> str:
        schema = extract_schema(self.db_instance)
        return self._format_schema(schema)

    @abstractmethod
    def _format_schema(self, schema: GraphSchema) -> str:
        pass

    def execute_query(self, query: str) -> DataFrame:
        result: QueryResult | None = None
        with self.db_instance.get_new_conn() as conn:
            try:
                result = conn.execute(query)
                df_result = result.get_as_df()
            except Exception as ex:
                logger.info(f"Failed to execute query {query}.")
                raise CypherQueryExecutionException(str(ex))
            finally:
                if result:
                    try:
                        result.close()
                    except Exception as ex:
                        logger.exception("Failed to close the query result.")

        return df_result


class KuzuLLMStrategy(LLMCypherDbStrategy):
    """Kuzu-specific implementation for LLM Cypher queries."""

    def _format_schema(self, schema: GraphSchema) -> str:
        node_lines = []
        for node_label, node_schema in schema["node_groups"].items():
            filtered_props = [
                prop for prop in node_schema["properties"]
                if self._should_include_property(prop["name"])
            ]

            current_obj = f"{node_label}"
            if filtered_props:
                props_str = ",".join([
                    f"({prop['name']}: {prop['type']})"
                    for prop in filtered_props
                ])
                current_obj += f": {props_str}"
            else:
                current_obj += ":"
            node_lines.append(current_obj)

        edge_lines = []
        for edge_label, edge_schema in schema["edge_groups"].items():
            from_label = edge_schema["from_group_label"]
            to_label = edge_schema["to_group_label"]
            edge_lines.append(f"(:{from_label})-[:{edge_label}]->(:{to_label})")

        node_section = "\n".join(node_lines)
        edge_section = "\n".join(edge_lines)

        formatted_schema = f"""
        NODE TYPES:
            {node_section}
        RELATIONSHIPS:
            {edge_section}
        """

        logger.debug(f"Kuzu Database Schema.\n{formatted_schema}")
        return formatted_schema

    @staticmethod
    def _should_include_property(prop_name: str) -> bool:
        if not prop_name.startswith("_dku"):
            return True
        return prop_name in ["source", "source type"]


class Neo4jLLMStrategy(LLMCypherDbStrategy):
    """Neo4j-specific implementation for LLM Cypher queries."""

    def _format_schema(self, schema: GraphSchema) -> str:
        node_lines = []
        for node_label, node_schema in schema["node_groups"].items():
            if node_schema["properties"]:
                sorted_props = sorted(node_schema["properties"], key=lambda p: p["name"])
                props = ", ".join([
                    f"{prop['name']}: {prop['type']}"
                    for prop in sorted_props
                ])
                node_lines.append(f":`{node_label}`: {{{props}}}")
            else:
                node_lines.append(f":`{node_label}`: {{}}")

        edge_lines = []
        for edge_label, edge_schema in schema["edge_groups"].items():
            if edge_schema["properties"]:
                sorted_edge_props = sorted(edge_schema["properties"], key=lambda p: p["name"])
                props = ", ".join([
                    f"{prop['name']}: {prop['type']}"
                    for prop in sorted_edge_props
                ])
                edge_lines.append(f"[::`{edge_label}` {{{props}}}]")
            else:
                edge_lines.append(f"[::`{edge_label}` {{None: None}}]")

        node_section = "\n".join(node_lines) if node_lines else "No node types found"
        edge_section = "\n".join(edge_lines) if edge_lines else "No relationship types found"

        formatted_schema = f"""
        NODE TYPES:
            {node_section}
        RELATIONSHIPS:
            {edge_section}
        """

        logger.debug(f"Neo4j Database Schema.\n{formatted_schema}")
        return formatted_schema
