"""Formula engine — DAG, topological sort, and evaluation.

Builds a dependency graph of formula fields, detects circular references,
and evaluates formulas in topological order.

Aggregate context (KTD3):
- ``SUM({f1})`` — f1 is an aggregate reference (entire column)
- ``{f1} + 1`` — f1 is a row reference (current record's value)
- ``{f1} + SUM({f2})`` — mixed: row f1 + column f2 sum

The engine distinguishes these by checking if a field reference appears
as a direct argument to an aggregate function.
"""

from __future__ import annotations

import ast
from collections import deque

from agentkit.bitable.formula.functions import AGGREGATE_FUNCTIONS, FUNCTION_REGISTRY
from agentkit.bitable.formula.parser import (
    FormulaDepthExceededError,
    FormulaParseError,
    FormulaSecurityError,
    UnknownFunctionError,
    evaluate_ast,
    parse_formula,
)


class CircularReferenceError(Exception):
    """Raised when formula fields form a circular dependency."""


class FormulaEngine:
    """Formula engine: parse, build DAG, detect cycles, evaluate.

    Usage::

        engine = FormulaEngine()
        # Register formula fields
        engine.add_formula(field_id="calc", formula="=SUM({src})", field_refs={"src"})
        # Evaluate for a specific record
        result = engine.evaluate(field_id="calc", field_values={"src": [1, 2, 3]})
    """

    def __init__(self) -> None:
        # field_id → (ast_tree, field_mapping, aggregate_refs, row_refs)
        self._formulas: dict[str, _FormulaEntry] = {}
        # DAG: field_id → set of field_ids it depends on
        self._dag: dict[str, set[str]] = {}

    def add_formula(self, field_id: str, formula: str) -> None:
        """Register a formula for a field.

        Raises:
            FormulaParseError: Syntax error.
            FormulaSecurityError: Disallowed AST node.
            UnknownFunctionError: Unregistered function.
            CircularReferenceError: Adding this formula creates a cycle.
        """
        tree, field_mapping, cross_table_mapping = parse_formula(
            formula, set(FUNCTION_REGISTRY.keys())
        )

        # Classify field refs into aggregate vs row context
        aggregate_refs, row_refs = _classify_refs(tree, field_mapping)

        entry = _FormulaEntry(
            tree=tree,
            field_mapping=field_mapping,
            aggregate_refs=aggregate_refs,
            row_refs=row_refs,
            cross_table_mapping=cross_table_mapping,
            formula=formula,
        )
        self._formulas[field_id] = entry
        # Update DAG: this field depends on all referenced fields
        # Cross-table refs depend on the relation field (local) + target field (foreign)
        cross_table_local_deps = {rel_id for rel_id, _ in cross_table_mapping.values()}
        self._dag[field_id] = aggregate_refs | row_refs | cross_table_local_deps

        # Check for cycles
        cycle = _detect_cycle(self._dag)
        if cycle:
            # Rollback
            del self._formulas[field_id]
            del self._dag[field_id]
            raise CircularReferenceError(f"Circular reference detected: {' → '.join(cycle)}")

    def remove_formula(self, field_id: str) -> None:
        """Remove a formula from the engine."""
        self._formulas.pop(field_id, None)
        self._dag.pop(field_id, None)
        # Remove edges pointing to this field
        for deps in self._dag.values():
            deps.discard(field_id)

    def get_dependencies(self, field_id: str) -> set[str]:
        """Get the set of field IDs that ``field_id`` depends on."""
        return self._dag.get(field_id, set()).copy()

    def get_cross_table_mapping(self, field_id: str) -> dict[str, tuple[str, str]] | None:
        """Get the cross-table reference mapping for a formula field.

        Returns ``{safe_name: (relation_field_id, target_field_id)}`` or
        ``None`` if the field is not registered or has no cross-table refs.
        """
        entry = self._formulas.get(field_id)
        if entry is None:
            return None
        return entry.cross_table_mapping or None

    def get_dependents(self, field_id: str) -> set[str]:
        """Get the set of formula field IDs that depend on ``field_id``."""
        return {fid for fid, deps in self._dag.items() if field_id in deps}

    def topological_order(self) -> list[str]:
        """Return all formula field IDs in topological order (Kahn's algorithm)."""
        return _topological_sort(self._dag)

    def evaluate(
        self,
        field_id: str,
        row_values: dict[str, object],
        column_values: dict[str, list[object]] | None = None,
        cross_table_values: dict[str, list[object]] | None = None,
    ) -> object:
        """Evaluate a formula field for a specific record.

        Args:
            field_id: The formula field to evaluate.
            row_values: Field ID → value for the current record (row context).
            column_values: Field ID → list of all values in that column
                (aggregate context). Required for aggregate references.
            cross_table_values: Safe-name → list of values from related
                records (cross-table context, U3). The service layer resolves
                these before calling evaluate.

        Returns:
            The computed value.

        Raises:
            KeyError: Field ID not registered.
            FormulaParseError: Field reference not found in values.
        """
        if field_id not in self._formulas:
            raise KeyError(f"Formula not registered: {field_id}")

        entry = self._formulas[field_id]
        column_values = column_values or {}
        cross_table_values = cross_table_values or {}

        # Build the field_values dict for the evaluator
        # Aggregate refs get column values (lists), row refs get row values (scalars)
        eval_values: dict[str, object] = {}

        # Map real field IDs to safe names
        for safe_name, real_id in entry.field_mapping.items():
            if real_id in entry.aggregate_refs:
                eval_values[safe_name] = column_values.get(real_id, [])
            else:
                eval_values[safe_name] = row_values.get(real_id)

        # Cross-table refs: the service layer pre-resolves these to value lists
        for safe_name in entry.cross_table_mapping:
            if safe_name in cross_table_values:
                eval_values[safe_name] = cross_table_values[safe_name]
            else:
                eval_values[safe_name] = []

        return evaluate_ast(entry.tree, eval_values, FUNCTION_REGISTRY)

    def evaluate_all_for_record(
        self,
        row_values: dict[str, object],
        column_values: dict[str, list[object]] | None = None,
    ) -> dict[str, object]:
        """Evaluate all registered formulas for a record.

        Returns a dict of field_id → computed value.
        Formulas are evaluated in topological order so that formula-to-formula
        dependencies are resolved correctly.

        List return values (from FILTER/SPLIT/LOOKUP) are stored as-is.
        Downstream formulas referencing a list-valued field receive the list.
        """
        results: dict[str, object] = {}
        column_values = column_values or {}

        for field_id in self.topological_order():
            # Include already-computed formula results in row_values
            merged_row = {**row_values, **results}
            try:
                results[field_id] = self.evaluate(field_id, merged_row, column_values)
            except (
                FormulaParseError,
                FormulaSecurityError,
                UnknownFunctionError,
                FormulaDepthExceededError,
                ZeroDivisionError,
                TypeError,
                ValueError,
            ) as e:
                results[field_id] = {"__error": str(e)}

        return results


# ── Internal data structures ──────────────────────────────


class _FormulaEntry:
    """Parsed formula metadata."""

    __slots__ = (
        "tree",
        "field_mapping",
        "aggregate_refs",
        "row_refs",
        "cross_table_mapping",
        "formula",
    )

    def __init__(
        self,
        tree: ast.Expression,
        field_mapping: dict[str, str],
        aggregate_refs: set[str],
        row_refs: set[str],
        cross_table_mapping: dict[str, tuple[str, str]],
        formula: str,
    ) -> None:
        self.tree = tree
        self.field_mapping = field_mapping
        self.aggregate_refs = aggregate_refs
        self.row_refs = row_refs
        self.cross_table_mapping = cross_table_mapping
        self.formula = formula


# ── DAG utilities ─────────────────────────────────────────


def _detect_cycle(dag: dict[str, set[str]]) -> list[str] | None:
    """Detect a cycle in the DAG using DFS. Returns the cycle path or None."""
    WHITE, GRAY, BLACK = 0, 1, 2
    color: dict[str, int] = {node: WHITE for node in dag}
    parent: dict[str, str | None] = {node: None for node in dag}

    def _dfs(node: str) -> list[str] | None:
        color[node] = GRAY
        for neighbor in dag.get(node, set()):
            if neighbor not in color:
                color[neighbor] = WHITE
                parent[neighbor] = None
            if color[neighbor] == GRAY:
                # Found cycle — reconstruct path
                cycle = [neighbor]
                current = node
                while current is not None and current != neighbor:
                    cycle.append(current)
                    current = parent.get(current)
                cycle.append(neighbor)
                cycle.reverse()
                return cycle
            if color[neighbor] == WHITE:
                parent[neighbor] = node
                result = _dfs(neighbor)
                if result is not None:
                    return result
        color[node] = BLACK
        return None

    for node in dag:
        if color.get(node, WHITE) == WHITE:
            result = _dfs(node)
            if result is not None:
                return result
    return None


def _topological_sort(dag: dict[str, set[str]]) -> list[str]:
    """Kahn's algorithm for topological sort."""
    # Build in-degree map
    in_degree: dict[str, int] = {node: 0 for node in dag}
    for node, deps in dag.items():
        for dep in deps:
            if dep in in_degree:
                in_degree[node] += 1

    # Start with nodes that have no dependencies
    queue = deque(node for node, degree in in_degree.items() if degree == 0)
    result: list[str] = []

    while queue:
        node = queue.popleft()
        result.append(node)
        # Find all nodes that depend on this node
        for other_node, deps in dag.items():
            if node in deps:
                in_degree[other_node] -= 1
                if in_degree[other_node] == 0 and other_node not in result:
                    queue.append(other_node)

    return result


def _classify_refs(
    tree: ast.Expression, field_mapping: dict[str, str]
) -> tuple[set[str], set[str]]:
    """Classify field references into aggregate (column) and row context.

    A field reference is aggregate if it appears as a direct argument to
    an aggregate function (SUM/AVG/COUNT/MIN/MAX). Otherwise it's row context.
    """
    aggregate_refs: set[str] = set()
    row_refs: set[str] = set()

    # Get all safe names → real field IDs
    safe_to_real = field_mapping

    class _Classifier(ast.NodeVisitor):
        def visit_Call(self, node: ast.Call) -> None:
            if isinstance(node.func, ast.Name) and node.func.id in AGGREGATE_FUNCTIONS:
                # Direct arguments to aggregate functions are column refs
                for arg in node.args:
                    if isinstance(arg, ast.Name) and arg.id in safe_to_real:
                        aggregate_refs.add(safe_to_real[arg.id])
                    else:
                        # Non-field args (e.g., literals) — visit normally
                        self.visit(arg)
            else:
                self.generic_visit(node)

        def visit_Name(self, node: ast.Name) -> None:
            if node.id in safe_to_real:
                real_id = safe_to_real[node.id]
                if real_id not in aggregate_refs:
                    row_refs.add(real_id)

    _Classifier().visit(tree)
    return aggregate_refs, row_refs