// direct.h — Computation graph with topological update
//
// Thread safety: Not thread-safe. External synchronization required.
// Behavior: dt_update does DFS from origin, evaluates nodes in reverse
//   order (topological). Uses a global generation counter to avoid
//   re-visiting nodes. The counter is u64 and overflows after ~1.8e19
//   calls, at which point reachable node generations are reset.
// Memory: Edges are heap-allocated per dt_connect. Use dt_clear_edges to
//   free all edges for a node.

#ifndef DIRECT_H
#define DIRECT_H

/* USER-DEFINED (define before including):
 *   typedef <type> direct_value_t;
 */

#include "type.h"

typedef struct edge {
  struct edge *next;
  struct node *node;
} edge_t;

typedef struct node {
  struct edge *edge;
  direct_value_t x;
  u64 visited_generation;
  direct_value_t (*function)(struct node *node);
} node_t;

int dt_connect(node_t *, node_t *);
int dt_connect_ok(node_t *, node_t *);
void dt_disconnect(node_t *, node_t *);
void dt_update(node_t *);
void dt_update_trace(node_t *);
void dt_clear_edges(node_t *);

direct_value_t dt_one(node_t *node);
direct_value_t dt_add(node_t *node);
direct_value_t dt_subtract(node_t *node);

#endif

#ifdef DT_IMPLEMENTATION
#ifndef DT_IMPLEMENTED
#define DT_IMPLEMENTED

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

// dt_connect() / dt_connect_ok()
//
// Add an edge from 'to' to 'from'. When dt_update is called on a node,
// it will recursively evaluate its connected dependents.
//
// to   - Destination node (depends on 'from')
// from - Source node
//
// return - 1 on success, 0 on failure

int dt_connect(node_t *to, node_t *from) {
  return dt_connect_ok(to, from);
}

int dt_connect_ok(node_t *to, node_t *from) {
  if (!to || !from)
    return 0;

  edge_t *edge = malloc(sizeof(edge_t));
  if (!edge)
    return 0;
  edge->node = from;
  edge->next = to->edge;
  to->edge = edge;
  return 1;
}

// dt_disconnect()
//
// Remove the edge from 'to' to 'from'

void dt_disconnect(node_t *to, node_t *from) {
  if (!to || !from)
    return;

  edge_t *edge = to->edge;
  edge_t *prev = NULL;

  while (edge) {
    if (edge->node == from) {
      if (prev) {
        prev->next = edge->next;
      } else {
        to->edge = edge->next;
      }
      free(edge);
      break;
    }
    prev = edge;
    edge = edge->next;
  }
}

static u64 _dt_generation = 1;

static int dt_grow(node_t ***s, node_t ***e, u32 *cap) {
  if (*cap > UINT32_MAX / 2) return 0;
  u32 nc = (*cap) * 2;
  size_t grow_sz = (size_t)nc * sizeof(node_t *);
  if (nc > 0 && grow_sz / nc != sizeof(node_t *)) return 0;
  node_t **ns = realloc(*s, grow_sz);
  if (!ns) return 0;
  node_t **ne = realloc(*e, grow_sz);
  if (!ne) {
    node_t **shrunk = realloc(*s, (size_t)(*cap) * sizeof(node_t *));
    if (!shrunk)
      shrunk = ns;
    else
      free(ns);
    *s = shrunk;
    return 0;
  }
  *s = ns;
  *e = ne;
  *cap = nc;
  return 1;
}

static void dt_update_impl(node_t *origin) {
  if (!origin || !origin->function)
    return;
  u32 cap = 64;
  node_t **dfs = malloc(cap * sizeof(node_t *));
  if (!dfs) return;
  node_t **eval = malloc(cap * sizeof(node_t *));
  if (!eval) { free(dfs); return; }
  u32 sp = 0, ep = 0;
  dfs[sp++] = origin;
  while (sp > 0) {
    node_t *c = dfs[--sp];
    if (!c || !c->function) continue;
    if (c->visited_generation == _dt_generation) continue;
    c->visited_generation = _dt_generation;
    if (!dt_grow(&dfs, &eval, &cap)) { free(dfs); free(eval); return; }
    eval[ep++] = c;
    edge_t *e = c->edge;
    while (e) {
      if (e->node && e->node->function && e->node->visited_generation != _dt_generation) {
        if (!dt_grow(&dfs, &eval, &cap)) { free(dfs); free(eval); return; }
        dfs[sp++] = e->node;
      }
      e = e->next;
    }
  }
  for (u32 i = ep; i > 0; i--)
    eval[i - 1]->x = eval[i - 1]->function(eval[i - 1]);
  free(dfs);
  free(eval);
}

static void dt_reset_generations(node_t *origin, node_t **dfs, u32 *dfs_sp,
                                   u32 cap, u32 *stack_cap) {
    (void)cap;
    if (*stack_cap == 0) return;
    node_t *c = origin;
    if (!c || !c->function) return;
    dfs[(*dfs_sp)++] = c;
    c->visited_generation = 0;
    while (*dfs_sp > 0) {
        node_t *n = dfs[--(*dfs_sp)];
        edge_t *e = n->edge;
        while (e) {
            if (e->node && e->node->function && e->node->visited_generation != 0) {
                if (*dfs_sp >= *stack_cap) {
                    u32 nc = (*stack_cap) * 2;
                    node_t **nd = realloc(dfs, (size_t)nc * sizeof(node_t *));
                    if (nd) { dfs = nd; *stack_cap = nc; }
                    else return;
                }
                dfs[(*dfs_sp)++] = e->node;
                e->node->visited_generation = 0;
            }
            e = e->next;
        }
    }
}

// dt_update()
//
// Evaluate the computation graph starting from origin. Performs DFS to
// collect all reachable nodes, then evaluates them in reverse (topological)
// order. Uses generation counters to skip already-evaluated nodes.

void dt_update(node_t *origin) {
  if (!origin || !origin->function)
    return;
  if (_dt_generation >= UINT64_MAX) {
    _dt_generation = 1;
    u32 cap = 64;
    u32 sp = 0;
    node_t **dfs = malloc(cap * sizeof(node_t *));
    if (dfs) {
        dt_reset_generations(origin, dfs, &sp, cap, &cap);
        free(dfs);
    }
  } else {
    _dt_generation++;
  }
  dt_update_impl(origin);
}

static void dt_update_trace_impl(node_t *origin) {
  if (!origin || !origin->function)
    return;
  u32 cap = 64;
  node_t **dfs = malloc(cap * sizeof(node_t *));
  if (!dfs) return;
  node_t **eval = malloc(cap * sizeof(node_t *));
  if (!eval) { free(dfs); return; }
  u32 sp = 0, ep = 0;
  dfs[sp++] = origin;
  while (sp > 0) {
    node_t *c = dfs[--sp];
    if (!c || !c->function) continue;
    if (c->visited_generation == _dt_generation) continue;
    c->visited_generation = _dt_generation;
    if (!dt_grow(&dfs, &eval, &cap)) { free(dfs); free(eval); return; }
    eval[ep++] = c;
    fprintf(stderr, ">: %p\n", (void *)c);
    edge_t *e = c->edge;
    while (e) {
      if (e->node && e->node->function && e->node->visited_generation != _dt_generation) {
        if (!dt_grow(&dfs, &eval, &cap)) { free(dfs); free(eval); return; }
        dfs[sp++] = e->node;
      }
      e = e->next;
    }
  }
  for (u32 i = ep; i > 0; i--) {
    node_t *n = eval[i - 1];
    n->x = n->function(n);
    fprintf(stderr, "<: %p\n", (void *)n);
  }
  free(dfs);
  free(eval);
}

// dt_update_trace()
//
// Same as dt_update, but prints each node's address to stderr as it is
// visited (">:") and evaluated ("<:") for debugging purposes.

void dt_update_trace(node_t *origin) {
  if (!origin || !origin->function)
    return;
  if (_dt_generation >= UINT64_MAX) {
    _dt_generation = 1;
    u32 cap = 64;
    u32 sp = 0;
    node_t **dfs = malloc(cap * sizeof(node_t *));
    if (dfs) {
        dt_reset_generations(origin, dfs, &sp, cap, &cap);
        free(dfs);
    }
  } else {
    _dt_generation++;
  }
  dt_update_trace_impl(origin);
}

// dt_one() - Constant node that always returns 1
direct_value_t dt_one(node_t *node) { (void)node; return 1; }

// dt_add() - Sums the values of all connected input nodes plus the node's own value
direct_value_t dt_add(node_t *node) {
  if (!node)
    return 0;
  edge_t *edge = node->edge;
  direct_value_t x = node->x;

  while (edge) {
    if (edge->node)
      x += edge->node->x;
    edge = edge->next;
  }

  return x;
}

// dt_subtract() - Subtracts all connected input node values from the node's own value
direct_value_t dt_subtract(node_t *node) {
  if (!node)
    return 0;
  edge_t *edge = node->edge;
  direct_value_t x = node->x;

  while (edge) {
    if (edge->node)
      x -= edge->node->x;
    edge = edge->next;
  }

  return x;
}

// dt_clear_edges()
//
// Free all edges connected to a node

void dt_clear_edges(node_t *node) {
  if (!node)
    return;
  edge_t *edge = node->edge;
  node->edge = NULL;

  while (edge) {
    edge_t *next = edge->next;
    free(edge);
    edge = next;
  }
}

#endif /* DT_IMPLEMENTED */
#endif /* DT_IMPLEMENTATION */