from __future__ import annotations

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

import networkx as nx
from neo4j import Record
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

logger = logging.getLogger(__name__)


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, records: list[Record]) -> None:
        """
        Initialize Neo4j query result wrapper.
        
        Args:
            records: List of Neo4j Record objects returned from query execution.
        """
        self.__records = records

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

        self.__extracted = False
    
    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 not rel.start_node or not rel.end_node:
            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.__records:
            for value in record.values():
                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.
        """
        if not self.__records:
            return DataFrame()

        data = [
            {key: self.__to_serialize_value(value) for key, value in record.items()}
            for record in self.__records
        ]      
  
        logger.debug(f"Converted {len(data)} records to DataFrame")
        return DataFrame(data)

    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
