backend/app/graph/graph_builder.py

import math

import pandas as pd
from pydantic import BaseModel

from app.config.constants import BATCH_SIZE
from app.graph.neo4j_client import Neo4jClient
from app.modeling.graph_schema import EdgeSchema, GraphSchema
from app.utils.logger import get_logger

logger = get_logger("graph.graph_builder")


class BuildReport(BaseModel):
    nodes_created: dict[str, int] = {}
    edges_created: dict[str, int] = {}
    errors: list[str] = []


class GraphBuilder:
    """Uses MERGE (not CREATE) for idempotent inserts. Uses id_field from NodeSchema as merge key."""

    def __init__(self, client: Neo4jClient, schema: GraphSchema) -> None:
        self.client = client
        self.schema = schema

    # ------------------------------------------------------------------
    # Nodes
    # ------------------------------------------------------------------

    def merge_nodes(self, label: str, df: pd.DataFrame, id_fields: list[str]) -> int:
        """MERGE nodes for a single entity type in batches.

        Supports composite keys: id_fields can be ["salesOrder", "salesOrderItem"].
        """
        total = 0
        rows = self._df_to_records(df)
        batches = self._chunk(rows, BATCH_SIZE)

        # Build MERGE key clause: {field1: row.field1, field2: row.field2}
        merge_props = ", ".join(f"`{f}`: row.`{f}`" for f in id_fields)
        query = (
            f"UNWIND $rows AS row "
            f"MERGE (n:`{label}` {{{merge_props}}}) "
            f"SET n += row"
        )

        for i, batch in enumerate(batches):
            try:
                self.client.execute_write(query, {"rows": batch})
                total += len(batch)
            except Exception as e:
                logger.error("Failed merging nodes %s batch %d: %s", label, i, e)
                raise

        logger.info("Merged %d nodes for %s", total, label)
        return total

    # ------------------------------------------------------------------
    # Edges
    # ------------------------------------------------------------------

    def merge_edges(self, edge: EdgeSchema, dataframes: dict[str, pd.DataFrame]) -> int:
        """MERGE edges between two node types based on join fields.

        Supports two modes:
        1. Direct FK: reads from from_node's DataFrame (original behaviour).
        2. Bridge table: when ``edge.source_entity`` is set, reads from the
           bridge table DataFrame and uses ``edge.bridge_target_field`` for
           target matching and ``edge.rel_properties`` for relationship props.
        """
        src_field, tgt_field = edge.join_on

        # --- Determine which DataFrame supplies the join data ---------------
        if edge.source_entity and edge.source_entity in dataframes:
            join_df = dataframes[edge.source_entity]
        else:
            join_df = self._find_df_for_label(edge.from_node, dataframes)

        if join_df is None:
            logger.warning("Skipping edge %s: no DataFrame found", edge.type)
            return 0

        # Column in join_df that holds the target node's key value
        if edge.source_entity:
            # Bridge table: use bridge_target_field or fall back to tgt_field
            bridge_tgt_col = edge.bridge_target_field or tgt_field
        else:
            # Direct FK: the value comes from src_field in the from_node's DF
            bridge_tgt_col = src_field

        # Validate required columns exist
        if src_field not in join_df.columns:
            logger.warning(
                "Skipping edge %s: source field '%s' not in DataFrame (has %s)",
                edge.type, src_field, list(join_df.columns)[:10],
            )
            return 0
        if bridge_tgt_col not in join_df.columns:
            logger.warning(
                "Skipping edge %s: target field '%s' not in DataFrame",
                edge.type, bridge_tgt_col,
            )
            return 0

        src_node = self.schema.get_node(edge.from_node)
        tgt_node = self.schema.get_node(edge.to_node)
        if not src_node or not tgt_node:
            return 0

        # --- Collect needed columns -----------------------------------------
        src_id_fields = src_node.id_fields
        needed_cols = list(dict.fromkeys(src_id_fields + [src_field, bridge_tgt_col]))

        # Add relationship property columns
        rel_prop_cols: list[str] = []
        if edge.rel_properties:
            for prop in edge.rel_properties:
                if prop in join_df.columns and prop not in needed_cols:
                    needed_cols.append(prop)
                    rel_prop_cols.append(prop)

        missing = [c for c in needed_cols if c not in join_df.columns]
        if missing:
            logger.warning("Skipping edge %s: missing columns %s", edge.type, missing)
            return 0

        pairs = (
            join_df[needed_cols]
            .dropna(subset=[src_field, bridge_tgt_col])
            .drop_duplicates()
        )
        records = self._df_to_records(pairs)

        # --- Build Cypher ---------------------------------------------------
        src_match = ", ".join(f"`{f}`: row.`{f}`" for f in src_id_fields)
        query = (
            f"UNWIND $rows AS row "
            f"MATCH (a:`{edge.from_node}` {{{src_match}}}) "
            f"MATCH (b:`{edge.to_node}` {{`{tgt_field}`: row.`{bridge_tgt_col}`}}) "
            f"MERGE (a)-[r:`{edge.type}`]->(b)"
        )

        # SET relationship properties
        if rel_prop_cols:
            set_parts = [f"r.`{p}` = row.`{p}`" for p in rel_prop_cols]
            query += " SET " + ", ".join(set_parts)

        total = 0
        batches = self._chunk(records, BATCH_SIZE)
        for i, batch in enumerate(batches):
            try:
                self.client.execute_write(query, {"rows": batch})
                total += len(batch)
            except Exception as e:
                logger.error("Failed merging edges %s batch %d: %s", edge.type, i, e)
                raise

        logger.info(
            "Merged edges [%s] (%s)->(%s): %d pairs processed",
            edge.type, edge.from_node, edge.to_node, total,
        )
        return total

    # ------------------------------------------------------------------
    # Full build
    # ------------------------------------------------------------------

    def build_all(self, dataframes: dict[str, pd.DataFrame]) -> BuildReport:
        """Build the entire graph: nodes first, then edges."""
        report = BuildReport()

        # 1. Create indexes for merge performance
        self._create_indexes()

        # 2. Merge all nodes
        for node in self.schema.nodes:
            df = self._find_df_for_label(node.label, dataframes)
            if df is None:
                report.errors.append(f"No DataFrame found for node {node.label}")
                continue
            try:
                count = self.merge_nodes(node.label, df, node.id_fields)
                report.nodes_created[node.label] = count
            except Exception as e:
                report.errors.append(f"Node {node.label}: {e}")

        # 3. Merge all edges
        for edge in self.schema.edges:
            try:
                count = self.merge_edges(edge, dataframes)
                report.edges_created[edge.type] = report.edges_created.get(edge.type, 0) + count
            except Exception as e:
                report.errors.append(f"Edge {edge.type}: {e}")

        # 4. Post-build sanity checks
        self._run_sanity_checks(report)

        logger.info(
            "Build complete: %d node types, %d edge types, %d errors",
            len(report.nodes_created), len(report.edges_created), len(report.errors),
        )
        return report

    def clear_graph(self) -> None:
        """Delete all nodes and relationships."""
        self.client.execute_write("MATCH (n) DETACH DELETE n")
        logger.info("Graph cleared")

    # ------------------------------------------------------------------
    # Post-build sanity checks
    # ------------------------------------------------------------------

    def _run_sanity_checks(self, report: BuildReport) -> None:
        """Run post-build graph sanity checks and append warnings to report."""
        # 1. Check for duplicate nodes per label
        for node in self.schema.nodes:
            id_str = ", ".join(f"n.`{f}`" for f in node.id_fields)
            query = (
                f"MATCH (n:`{node.label}`) "
                f"WITH {id_str}, count(*) AS cnt "
                f"WHERE cnt > 1 "
                f"RETURN {id_str}, cnt ORDER BY cnt DESC LIMIT 5"
            )
            try:
                results = self.client.execute(query)
                if results:
                    report.errors.append(f"Duplicate nodes in {node.label}: {results}")
            except Exception as e:
                logger.debug("Sanity check skipped for %s: %s", node.label, e)

        # 2. Check for extreme high-degree nodes (star anomaly)
        try:
            query = (
                "MATCH (n) "
                "WITH n, size([(n)--() | 1]) AS degree "
                "WHERE degree > 200 "
                "RETURN labels(n)[0] AS label, count(n) AS cnt"
            )
            results = self.client.execute(query)
            if results:
                report.errors.append(f"Star anomaly detected: {results}")
        except Exception as e:
            logger.debug("Star anomaly check skipped: %s", e)

        # 3. Check for isolated nodes
        try:
            query = (
                "MATCH (n) WHERE NOT (n)--() "
                "RETURN labels(n)[0] AS label, count(n) AS cnt"
            )
            results = self.client.execute(query)
            for r in results:
                logger.warning("Isolated nodes: %s = %s", r.get("label"), r.get("cnt"))
        except Exception as e:
            logger.debug("Isolation check skipped: %s", e)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _create_indexes(self) -> None:
        """Create indexes on id_field(s) for each node label to speed up MERGE."""
        for node in self.schema.nodes:
            fields = node.id_fields
            index_name = f"idx_{node.label}_{'_'.join(fields)}"
            props = ", ".join(f"n.`{f}`" for f in fields)
            query = (
                f"CREATE INDEX {index_name} IF NOT EXISTS "
                f"FOR (n:`{node.label}`) ON ({props})"
            )
            try:
                self.client.execute_write(query)
            except Exception as e:
                logger.debug("Index creation note for %s: %s", node.label, e)

    def _find_df_for_label(self, label: str, dataframes: dict[str, pd.DataFrame]) -> pd.DataFrame | None:
        """Find the DataFrame matching a node label, using source_entity when available."""
        # Check source_entity from schema first (supports label != folder name)
        for node in self.schema.nodes:
            if node.label == label and node.source_entity and node.source_entity in dataframes:
                return dataframes[node.source_entity]
        # Fallback: original sanitize_label matching
        for entity_name, df in dataframes.items():
            from app.utils.helpers import sanitize_label
            if sanitize_label(entity_name) == label:
                return df
        return None

    @staticmethod
    def _df_to_records(df: pd.DataFrame) -> list[dict]:
        """Convert DataFrame to list of dicts, stringifying values and dropping nested objects."""
        records = []
        for row in df.to_dict("records"):
            clean = {}
            for k, v in row.items():
                if isinstance(v, dict):
                    continue  # skip nested objects (e.g. creationTime)
                if v is None or (isinstance(v, float) and math.isnan(v)):
                    continue  # skip nulls
                clean[k] = str(v)
            records.append(clean)
        return records

    @staticmethod
    def _chunk(items: list, size: int) -> list[list]:
        return [items[i : i + size] for i in range(0, len(items), size)]