"""Bulk node and edge loading from DataFrames into AGE graphs."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from piggie.exceptions import BulkLoadError

if TYPE_CHECKING:
    import pandas as pd

    from piggie.driver import PiggieDriver

logger = logging.getLogger(__name__)

_DEFAULT_BATCH = 1000
_MIN_BATCH = 1


def _cypher_value(val: Any) -> str:
    """Format a Python value as a Cypher literal string for mogrify."""
    if val is None:
        return "NULL"
    if isinstance(val, bool):
        return "true" if val else "false"
    if isinstance(val, int):
        return str(val)
    if isinstance(val, float):
        return str(val)
    # Everything else → string, escaped for Cypher
    s = str(val).replace("\\", "\\\\").replace("'", "\\'")
    return f"'{s}'"


def _props_fragment(row: dict[str, Any], cols: list[str]) -> str:
    """Build a Cypher property map literal: {key1: val1, key2: val2}."""
    parts = []
    for col in cols:
        v = row.get(col)
        if v is not None:
            parts.append(f"{col}: {_cypher_value(v)}")
    return "{" + ", ".join(parts) + "}"


def load_nodes(
    driver: PiggieDriver,
    df: pd.DataFrame,
    graph: str,
    label: str,
    *,
    id_col: str | None = None,
    property_cols: list[str] | None = None,
    batch_size: int = _DEFAULT_BATCH,
) -> int:
    """Load nodes from a DataFrame into an AGE graph.

    Dual-mode:
    - CREATE (default): Fast append. Every row becomes a new node.
    - MERGE (when ``id_col`` is set): Upsert. Matches on ``id_col``,
      sets remaining properties.

    Args:
        driver: Connected PiggieDriver.
        df: DataFrame with node data.
        graph: Target graph name.
        label: Vertex label.
        id_col: Column to use as merge key. Enables MERGE mode.
        property_cols: Columns to load as properties. Defaults to all columns.
        batch_size: Rows per transaction batch.

    Returns:
        Number of nodes loaded.
    """
    if df.empty:
        return 0

    cols = property_cols or [c for c in df.columns]
    if id_col and id_col not in cols:
        cols = [id_col] + cols

    total = 0
    rows = df.to_dict("records")

    for start in range(0, len(rows), batch_size):
        batch = rows[start : start + batch_size]
        stmts: list[str] = []

        for row in batch:
            if id_col:
                # MERGE mode
                id_val = _cypher_value(row[id_col])
                set_cols = [c for c in cols if c != id_col]
                set_parts = ", ".join(
                    f"n.{c} = {_cypher_value(row.get(c))}" for c in set_cols if row.get(c) is not None
                )
                stmt = f"MERGE (n:{label} {{{id_col}: {id_val}}})"
                if set_parts:
                    stmt += f" SET {set_parts}"
                stmts.append(stmt)
            else:
                # CREATE mode
                props = _props_fragment(row, cols)
                stmts.append(f"CREATE (:{label} {props})")

        # Execute batch — one Cypher statement at a time since AGE
        # doesn't support multi-statement Cypher in a single call
        try:
            for stmt in stmts:
                driver.cypher(stmt, graph=graph)
            driver.commit()
            total += len(batch)
        except Exception as exc:
            driver.rollback()
            # Adaptive: if batch is larger than 1, try smaller batches
            if batch_size > _MIN_BATCH:
                new_size = max(_MIN_BATCH, batch_size // 2)
                logger.warning(
                    "Batch of %d failed, retrying with batch_size=%d: %s",
                    batch_size, new_size, exc,
                )
                remaining_df = df.iloc[start:]
                total += load_nodes(
                    driver, remaining_df, graph, label,
                    id_col=id_col, property_cols=property_cols,
                    batch_size=new_size,
                )
                return total
            raise BulkLoadError(f"Failed to load nodes: {exc}") from exc

    return total


def load_edges(
    driver: PiggieDriver,
    df: pd.DataFrame,
    graph: str,
    rel_type: str,
    source_col: str,
    source_label: str,
    target_col: str,
    target_label: str,
    *,
    source_key: str | None = None,
    target_key: str | None = None,
    property_cols: list[str] | None = None,
    batch_size: int = _DEFAULT_BATCH,
) -> int:
    """Load edges from a DataFrame into an AGE graph.

    Each row creates an edge between existing nodes. Source and target
    nodes are matched by their node property values.

    Args:
        driver: Connected PiggieDriver.
        df: DataFrame with edge data.
        graph: Target graph name.
        rel_type: Edge label.
        source_col: DataFrame column containing source node identifiers.
        source_label: Label of source nodes.
        target_col: DataFrame column containing target node identifiers.
        target_label: Label of target nodes.
        source_key: Node property to match source nodes on. Defaults to source_col.
        target_key: Node property to match target nodes on. Defaults to target_col.
        property_cols: Additional columns to store as edge properties.
        batch_size: Rows per transaction batch.

    Returns:
        Number of edges loaded.
    """
    if df.empty:
        return 0

    src_key = source_key or source_col
    tgt_key = target_key or target_col
    prop_cols = property_cols or []
    total = 0
    rows = df.to_dict("records")

    for start in range(0, len(rows), batch_size):
        batch = rows[start : start + batch_size]
        stmts: list[str] = []

        for row in batch:
            src_val = _cypher_value(row[source_col])
            tgt_val = _cypher_value(row[target_col])

            if prop_cols:
                props = _props_fragment(row, prop_cols)
                stmt = (
                    f"MATCH (a:{source_label} {{{src_key}: {src_val}}}), "
                    f"(b:{target_label} {{{tgt_key}: {tgt_val}}}) "
                    f"CREATE (a)-[:{rel_type} {props}]->(b)"
                )
            else:
                stmt = (
                    f"MATCH (a:{source_label} {{{src_key}: {src_val}}}), "
                    f"(b:{target_label} {{{tgt_key}: {tgt_val}}}) "
                    f"CREATE (a)-[:{rel_type}]->(b)"
                )
            stmts.append(stmt)

        try:
            for stmt in stmts:
                driver.cypher(stmt, graph=graph)
            driver.commit()
            total += len(batch)
        except Exception as exc:
            driver.rollback()
            if batch_size > _MIN_BATCH:
                new_size = max(_MIN_BATCH, batch_size // 2)
                logger.warning(
                    "Edge batch of %d failed, retrying with batch_size=%d: %s",
                    batch_size, new_size, exc,
                )
                remaining_df = df.iloc[start:]
                total += load_edges(
                    driver, remaining_df, graph, rel_type,
                    source_col, source_label, target_col, target_label,
                    source_key=source_key, target_key=target_key,
                    property_cols=property_cols, batch_size=new_size,
                )
                return total
            raise BulkLoadError(f"Failed to load edges: {exc}") from exc

    return total
