from __future__ import annotations

import logging
from datetime import date, datetime, time
from typing import Any, cast

import networkx as nx
from neo4j import Result
from neo4j.graph import Node, Path, Relationship
from neo4j.spatial import Point
from neo4j.time import Date, DateTime, Duration, Time
from networkx import MultiDiGraph
from pandas import DataFrame, isna

from graph.connections.query_result import NodeData, QueryResult, RelationshipData, TableData
from graph.ingestion.dataiku.data_types import GraphDataikuDataType, NativeDataikuDataType

logger = logging.getLogger(__name__)

PANDAS_DKU_TYPES_MAP: dict[str, NativeDataikuDataType] = {
    "int8": "tinyint",
    "int16": "smallint",
    "int32": "int",
    "int64": "bigint",
    "Int64": "bigint",
    "int": "int",
    "float32": "float",
    "float64": "double",
    "float": "float",
    "object": "string",
    "bool": "boolean",
    "boolean": "boolean",
}


class Neo4jQueryResult(QueryResult):
    """
    Wrapper for Neo4j query results providing conversion to DataFrame and NetworkX graphs.

    This class handles Neo4j-specific graph structures (Node, Relationship, Path) and
    converts them to standard data formats.
    """

    def __init__(self, result: Result) -> None:
        """
        Initialize Neo4j query result wrapper.

        Args:
            result: Neo4j Result object
        """
        self.__df = result.to_df()
        self.__schema: dict[str, GraphDataikuDataType] = Neo4jQueryResult.__extract_schema__(self.__df)

        self.__nodes: dict[str, NodeData] = {}
        self.__relationships: dict[str, RelationshipData] = {}

        self.__extracted = False

    @staticmethod
    def __extract_schema__(df: DataFrame) -> dict[str, GraphDataikuDataType]:
        schema: dict[str, GraphDataikuDataType] = {}

        for col in df.columns:
            # Get the first valid value to check type
            non_empty_rows = df[col].dropna()
            sample_value = non_empty_rows.iloc[0] if not non_empty_rows.empty else None

            schema[col] = Neo4jQueryResult.__infer_dataiku_type__(sample_value, str(df[col].dtype))

        return schema

    @staticmethod
    def __infer_dataiku_type__(sample_value: Any, pandas_dtype: str) -> GraphDataikuDataType:
        inferred_type: GraphDataikuDataType

        if isinstance(sample_value, Node):
            inferred_type = "node"
        elif isinstance(sample_value, Relationship):
            inferred_type = "edge"
        elif isinstance(sample_value, Date):
            inferred_type = "dateonly"
        elif isinstance(sample_value, DateTime):
            inferred_type = "date"
        elif isinstance(sample_value, date):
            # Specific python type, datetime is a subclass of date.
            inferred_type = "date" if isinstance(sample_value, datetime) else "dateonly"
        elif isinstance(sample_value, (Time, time, Duration)):
            inferred_type = "string"
        elif isinstance(sample_value, Path):
            inferred_type = "path"
        elif isinstance(sample_value, list):
            inferred_type = "array"
        elif isinstance(sample_value, Point):
            inferred_type = "geopoint"
        elif isinstance(sample_value, dict):
            inferred_type = "object"
        else:
            inferred_type = PANDAS_DKU_TYPES_MAP.get(
                pandas_dtype, "object"
            )  # Default to object if we cannot infer a more specific type

        return inferred_type

    def __to_serializable_property(self, property: Any) -> Any:
        """
        Serialize a property value from a Neo4j record.

        This method handles different data types that can be returned as
        properties from Neo4j, converting them into Python-friendly formats
        suitable for JSON serialization or other uses.

        Args:
            property: The property value to serialize.  It can be of any type:
                primitive types, lists, dictionaries, or Neo4j-specific types
                like dates, times, datetimes, durations, and points.

        Returns:
            A serialized representation of the property.  Neo4j date and time
            objects are converted to ISO format strings.  Neo4j spatial `Point`
            objects are converted to dictionaries with 'x', 'y', and 'srid' keys.
            `NaN` values are converted to `None`. Lists and dictionaries are
            processed recursively.
        """

        result = property

        if isinstance(property, list):
            result = [self.__to_serializable_property(item) for item in property]

        elif isinstance(property, dict):
            result = {k: self.__to_serializable_property(v) for k, v in property.items()}

        elif isinstance(property, (date, time, datetime, Date, Time, DateTime)):
            result = property.isoformat()

        elif isinstance(property, Duration):
            result = str(property)

        elif isinstance(property, Point):
            res = {"x": property.x, "y": property.y, "srid": property.srid}
            if hasattr(property, "z") and property.z is not None:
                res["z"] = property.z
            result = res

        elif isna(property):
            result = None

        return result

    def __to_serialize_node(self, node: Node) -> tuple[str, NodeData]:
        """
        Serialize a Neo4j Node to a dictionary.

        Args:
            node: Neo4j Node object.

        Returns:
            Dictionary with node properties and metadata.
        """
        labels = list(node.labels)
        label = labels[0] if labels else ""
        dku_grp_def_id = node.get("_dku_grp_def_id", label)
        dku_reserved_user_defined_pk_prop = node.get("_dku_reserved_user_defined_pk_prop", "id")

        data = NodeData(
            id=node.element_id,
            _id=node.element_id,
            _label=label,
            _dku_grp_def_id=dku_grp_def_id,
            _dku_reserved_user_defined_pk_prop=dku_reserved_user_defined_pk_prop,
        )

        data.update(**{k: self.__to_serializable_property(v) for k, v in node.items()})  # type: ignore

        return node.element_id, cast(NodeData, data)

    def __to_serialize_relationship(self, rel: Relationship) -> tuple[str, RelationshipData] | tuple[None, None]:
        """
        Serialize a Neo4j Relationship to a dictionary.

        Args:
            rel: Neo4j Relationship object.

        Returns:
            Dictionary with relationship properties and metadata, or None if invalid.
        """
        if rel.start_node is None or rel.end_node is None:
            logger.warning(f"Relationship {rel.element_id} missing start or end node. Skipping.")
            return None, None

        data = RelationshipData(
            _src=rel.start_node.element_id,
            _dst=rel.end_node.element_id,
            _label=rel.type,
            _dku_grp_def_id=rel.get("_dku_grp_def_id", rel.type),
        )

        data.update(**{k: self.__to_serializable_property(v) for k, v in rel.items()})  # type: ignore

        return rel.element_id, cast(RelationshipData, data)

    def __to_serialize_value(self, value: Any) -> Any:
        """
        Serialize a value from a Neo4j record.

        Args:
            value: Value to serialize (can be Node, Relationship, Path, or primitive).

        Returns:
            Serialized representation of the value.
            For Path objects, returns all nodes and relationships as a flat list.
        """
        if isinstance(value, Node):
            _, node_data = self.__to_serialize_node(value)
            return node_data
        elif isinstance(value, Relationship):
            _, rel_data = self.__to_serialize_relationship(value)
            return rel_data
        elif isinstance(value, Path):
            elements: list[NodeData | RelationshipData] = []
            for node in value.nodes:
                _, node_data = self.__to_serialize_node(node)
                elements.append(node_data)
            for rel in value.relationships:
                _, rel_data = self.__to_serialize_relationship(rel)
                if rel_data:
                    elements.append(rel_data)
            return elements
        else:
            return self.__to_serializable_property(value)

    def __extract_graph_elements(self) -> None:
        """
        Extract and cache nodes and relationships from records.

        This method processes all records and extracts Neo4j graph elements,
        storing them for efficient graph construction.
        """
        for _, record in self.__df.iterrows():
            for _, value in record.items():
                if isinstance(value, Node):
                    node_id, node_data = self.__to_serialize_node(value)
                    self.__nodes[node_id] = node_data

                elif isinstance(value, Relationship):
                    rel_id, rel_data = self.__to_serialize_relationship(value)
                    if rel_id is not None and rel_data is not None:
                        self.__relationships[rel_id] = rel_data

                elif isinstance(value, Path):
                    for node in value.nodes:
                        node_id, node_data = self.__to_serialize_node(node)
                        self.__nodes[node_id] = node_data
                    for rel in value.relationships:
                        rel_id, rel_data = self.__to_serialize_relationship(rel)
                        if rel_id is not None and rel_data is not None:
                            self.__relationships[rel_id] = rel_data

        self.__extracted = True

    def get_as_df(self) -> DataFrame:
        """
        Convert query results to a pandas DataFrame.

        Returns:
            DataFrame where each row represents a record and columns are the record keys.
            Neo4j graph structures are serialized to nested dictionaries.
        """
        data = [
            {key: self.__to_serialize_value(value) for key, value in record.items()}
            for _, record in self.__df.iterrows()
        ]

        logger.debug(f"Converted {len(data)} records to DataFrame")
        return DataFrame(data)

    def get_as_table(self) -> TableData:
        return TableData(df=self.get_as_df(), schema=self.__schema)

    def get_as_networkx(self) -> MultiDiGraph:
        """
        Convert query results to a NetworkX MultiDiGraph.

        Returns:
            MultiDiGraph containing all nodes and relationships from the query results.
            Node and edge attributes include all properties from Neo4j.
        """
        if not self.__extracted:  # Only extract once
            self.__extract_graph_elements()

        graph: nx.MultiDiGraph = nx.MultiDiGraph()

        # Add all nodes with their attributes
        for node_id, node_data in self.__nodes.items():
            graph.add_node(node_id, **node_data)

        # Add all edges with their attributes
        for rel_data in self.__relationships.values():
            src_id = rel_data["_src"]
            dst_id = rel_data["_dst"]

            # Only add edges if both nodes exist
            if src_id in self.__nodes and dst_id in self.__nodes:
                graph.add_edge(src_id, dst_id, **rel_data)
            else:
                logger.warning(f"Skipping relationship {src_id}-{dst_id}: missing node(s) {src_id} or {dst_id}")

        logger.debug(f"Created NetworkX graph with {graph.number_of_nodes()} nodes and {graph.number_of_edges()} edges")
        return graph

    def close(self) -> None:
        pass
