import logging
from typing import Dict, List, Set, Tuple

import neo4j.graph

from graph.connections.db_connection import DbConnection
from graph.connections.db_instance import DbInstance
from graph.exploration.schema.schema_extractor import (
    EdgePropertySchema,
    EdgeTableSchema,
    GraphSchema,
    NodePropertySchema,
    NodeTableSchema,
)

logger = logging.getLogger(__name__)


def _is_node_primary_key_constraint(record: object) -> bool:
    # Older Neo4j versions exposed node PK/uniqueness constraints with legacy
    # `type` values like `UNIQUENESS` and `NODE_KEY`, so filtering directly on
    # those values used to work. Newer versions can emit more specific names
    # such as `NODE_PROPERTY_UNIQUENESS` and include `entityType`, which means
    # the old filter misses valid node constraints and PKs stop being detected.
    entity_type = str(record.get("entityType", "") or "").upper()  # type: ignore[attr-defined]
    constraint_type = str(record.get("type", "") or "").upper()  # type: ignore[attr-defined]

    if entity_type and entity_type != "NODE":
        return False

    if constraint_type == "NODE_KEY":
        return True

    if "RELATIONSHIP" in constraint_type:
        return False

    return "UNIQUENESS" in constraint_type


def _get_pk_prop_to_node_label_map(connection: DbConnection) -> Set[Tuple[str, str]]:
    constraints_df = connection.execute("SHOW CONSTRAINTS").get_as_df()

    label_to_pk_prop_name: Set[Tuple[str, str]] = set()
    for _, record in constraints_df.iterrows():
        if not _is_node_primary_key_constraint(record):
            continue

        labels = record["labelsOrTypes"]
        props = record["properties"]

        if not labels or not props:
            continue

        for label in labels:
            for prop in props:
                label_to_pk_prop_name.add((label, prop))

    return label_to_pk_prop_name


def _get_nodes_table_schema(connection: DbConnection) -> Dict[str, NodeTableSchema]:
    """
    Limitation: multiple types for a given property is not supported. Default to first type in the list.
    """
    label_to_pk_prop_name = _get_pk_prop_to_node_label_map(connection)

    node_props_df = connection.execute(
        """
        CALL db.schema.nodeTypeProperties()
        YIELD nodeLabels, propertyName, propertyTypes, mandatory
        RETURN nodeLabels, propertyName, propertyTypes, mandatory
        ORDER BY nodeLabels;
        """
    ).get_as_df()

    node_labels_prop_map: Dict[str, List[NodePropertySchema]] = {}

    for _, prop_row in node_props_df.iterrows():
        node_labels = prop_row["nodeLabels"]
        prop_name = prop_row["propertyName"]
        prop_types = prop_row["propertyTypes"]

        if len(prop_types) > 1:
            logger.warning(
                "Multiple types for a given node property is not supported. Defaulting to first type '%s' for property '%s' on node labels '%s'. Types: %s.",
                prop_types[0],
                prop_name,
                node_labels,
                prop_types,
            )
        elif not prop_name or not prop_types:
            logger.warning(
                "Skipping property on node labels '%s' due to missing info. Property name: %s, Property types: %s.",
                node_labels,
                prop_name,
                prop_types,
            )
            continue

        for node_label in node_labels:
            if node_label not in node_labels_prop_map:
                node_labels_prop_map[node_label] = []

            if prop_name:
                node_labels_prop_map[node_label].append(
                    {
                        "name": prop_name,
                        "type": prop_types[0],
                        "is_primary_key": (node_label, prop_name) in label_to_pk_prop_name,
                    }
                )

    nodes: Dict[str, NodeTableSchema] = {}
    for node_label, properties in node_labels_prop_map.items():
        nodes[node_label] = {
            "label": node_label,
            "properties": properties,
        }

    return nodes


def _get_props_to_rel_type_map(connection: DbConnection) -> Dict[str, List[EdgePropertySchema]]:
    rel_props_df = connection.execute(
        """
            CALL db.schema.relTypeProperties()
            YIELD relType, propertyName, propertyTypes
            RETURN relType, propertyName, propertyTypes;
            """
    ).get_as_df()
    rel_type_props_map: Dict[str, List[EdgePropertySchema]] = {}
    for _, record in rel_props_df.iterrows():
        rel_type = record["relType"].lstrip(":").strip("`")
        prop_name = record["propertyName"]
        prop_types = record["propertyTypes"]

        if not prop_name or not prop_types:
            logger.warning(
                "Skipping property on edge type '%s' due to missing info. Property name: %s, Property types: %s.",
                rel_type,
                prop_name,
                prop_types,
            )
            if rel_type not in rel_type_props_map:
                rel_type_props_map[rel_type] = []
            continue

        if len(prop_types) > 1:
            logger.warning(
                "Multiple types for a given edge property is not supported. Defaulting to first type '%s' for property '%s' on edge type '%s'. Types: %s.",
                prop_types[0],
                prop_name,
                rel_type,
                prop_types,
            )

        if rel_type not in rel_type_props_map:
            rel_type_props_map[rel_type] = []

        rel_type_props_map[rel_type].append({"name": prop_name, "type": prop_types[0]})

    return rel_type_props_map


def _get_edges_table_schema(connection: DbConnection) -> Dict[str, EdgeTableSchema]:
    """
    Limitation: relationships between multiple set of nodes are not supported.
    """
    rel_type_props_map = _get_props_to_rel_type_map(connection)

    # Topology extraction to get from/to node labels for relationships.
    topology_df = connection.execute("CALL db.schema.visualization()").get_as_df()

    edges: Dict[str, EdgeTableSchema] = {}
    if not topology_df.empty:
        if len(topology_df) > 1:
            logger.warning(
                "Expected a single row from 'CALL db.schema.visualization()' but got %d rows. Using the first row only.",
                len(topology_df),
            )

        schema_nodes_list: List[neo4j.graph.Node] = topology_df.iloc[0]["nodes"]
        schema_rels_list: List[neo4j.graph.Relationship] = topology_df.iloc[0]["relationships"]

        node_id_to_label = {n.id: n.labels for n in schema_nodes_list}

        for rel in schema_rels_list:
            rel_type = rel.type

            if rel.start_node is None or rel.end_node is None:
                raise ValueError(f"Relationship of type '{rel_type}' is missing start or end node.")

            from_labels = list(node_id_to_label.get(rel.start_node.id, []))
            to_labels = list(node_id_to_label.get(rel.end_node.id, []))

            if not from_labels or not to_labels:
                raise ValueError(
                    f"Could not find from/to labels for relationship type '{rel_type}' with start node ID '{rel.start_node.id}' and end node ID '{rel.end_node.id}'."
                )

            if len(from_labels) > 1 or len(to_labels) > 1:
                raise NotImplementedError(
                    f"Multiple sets of node labels for a relationship are not supported. Expected a single label for from/to nodes but got multiple. From labels: {from_labels}, To labels: {to_labels}."
                )

            if rel_type in edges:
                raise NotImplementedError(
                    f"Multiple sets of node labels for a relationship is not supported. It was {edges[rel_type]}. Now attempting to add from labels: {from_labels}, to labels: {to_labels}."
                )

            edges[rel_type] = {
                "label": rel_type,
                "from_group_label": from_labels[0],
                "to_group_label": to_labels[0],
                "properties": rel_type_props_map.get(rel_type, []),
            }

    return edges


def extract_schema(db_instance: DbInstance) -> GraphSchema:
    with db_instance.get_new_conn() as connection:
        nodes: Dict[str, NodeTableSchema] = _get_nodes_table_schema(connection)
        edges: Dict[str, EdgeTableSchema] = _get_edges_table_schema(connection)

    return {"node_groups": nodes, "edge_groups": edges}
