import { BaseNode } from './nodes/baseNode';
import { ReflectionNode } from './nodes/reflectionNode';
import { RewriteNode, RewriteFunction } from './nodes/rewriteNode';

export type Edge = {
  from: string;
  out: string;
  to: string;
  in: string;
};

export class Graph {
  private nodes: Map<string, BaseNode>;
  private edges: Edge[];
  private nodeCounter: number;

  constructor() {
    this.nodes = new Map();
    this.edges = [];
    this.nodeCounter = 0;
  }

  private generateId(): string {
    return `node_${this.nodeCounter++}`;
  }

  /**
   * Creates a node of the specified type.
   * @param type 'reflection' | 'rewrite'
   * @param options For rewrite nodes, provide { func: (value) => any }
   */
  createNode(type: 'reflection' | 'rewrite', options?: any): BaseNode {
    const id = this.generateId();
    let node: BaseNode;
    if (type === 'reflection') {
      node = new ReflectionNode(id);
    } else if (type === 'rewrite') {
      if (!options || typeof options.func !== 'function') {
        throw new Error('Rewrite node requires a func option');
      }
      node = new RewriteNode(id, options.func);
    } else {
      throw new Error(`Unknown node type: ${type}`);
    }
    this.nodes.set(id, node);
    return node;
  }

  addNode(node: BaseNode): void {
    if (this.nodes.has(node.id)) {
      throw new Error(`Node with id ${node.id} already exists`);
    }
    this.nodes.set(node.id, node);
  }

  addEdge(from: string, out: string, to: string, inKey: string): void {
    if (!this.nodes.has(from) || !this.nodes.has(to)) {
      throw new Error('Both nodes must exist to add an edge');
    }
    this.edges.push({ from, out, to, in: inKey });
  }

  /**
   * Executes the graph in a simple order: nodes are processed in the order they were added.
   * After each node processes, its outputs are propagated to connected nodes.
   */
  run(): void {
    for (const node of this.nodes.values()) {
      node.process();
      for (const edge of this.edges.filter(e => e.from === node.id)) {
        const target = this.nodes.get(edge.to);
        if (!target) continue;
        const value = node.outputs.get(edge.out);
        target.inputs.set(edge.in, value);
      }
    }
  }

  getNode(id: string): BaseNode | undefined {
    return this.nodes.get(id);
  }
}