class LittleGhost::Graph

# File lib/little_ghost/graph.rb, line 164
def edge(from, to, input: nil, max_concurrency: nil, **options, &condition)
  condition = extract_condition(options, condition)
  validate_callable!(input, "edge input mapper")
  from = normalize_endpoint(from, "edge source")
  to = normalize_endpoint(to, "edge target")
  if from.is_a?(Array) && to.is_a?(Array)
    raise ArgumentError, "graph edges cannot use arrays for both source and target"
  end
  if from.is_a?(Array) && condition
    raise ArgumentError, "fan-in edges cannot be conditional"
  end
  unless max_concurrency.nil?
    raise ArgumentError, "max_concurrency is only valid for a fan-out edge" unless to.is_a?(Array)

    max_concurrency = normalize_max_concurrency(max_concurrency)
  end
  declaration = Edge.new(
    from:,
    to:,
    condition:,
    input_mapper: input,
    max_concurrency:
  )
  self.graph_edges_value = (graph_edges_value + [declaration]).freeze
  declaration
end

# File lib/little_ghost/graph.rb, line 195
def error_edge(from, to, on:, input: nil)
  errors = Array(on)
  unless errors.any? && errors.all? { |error| error.is_a?(Class) && error <= Exception }
    raise ArgumentError, "error edge on: must contain exception classes"
  end
  validate_callable!(input, "error edge input mapper")
  declaration = ErrorEdge.new(
    from: normalize_node_name(from),
    to: normalize_node_name(to),
    errors: errors.freeze,
    input_mapper: input
  )
  self.graph_error_edges_value = (graph_error_edges_value + [declaration]).freeze
  declaration
end

# File lib/little_ghost/graph.rb, line 212
def finish(name = nil)
  return graph_finish_value if name.nil?

  self.graph_finish_value = normalize_node_name(name)
end

# File lib/little_ghost/graph.rb, line 231
def max_concurrency(value = nil)
  return graph_max_concurrency_value if value.nil?

  self.graph_max_concurrency_value = normalize_max_concurrency(value)
end

# File lib/little_ghost/graph.rb, line 219
def max_steps(value = nil)
  return graph_max_steps_value if value.nil?

  value = Integer(value)
  raise ArgumentError, "max_steps must be at least 1" if value < 1

  self.graph_max_steps_value = value
end

# File lib/little_ghost/graph.rb, line 123
def node(name, assembly, timeout: nil, retries: 0, retry_on: nil, retry_delay: 0,
  history: false, context: false, input: nil)
  name = normalize_node_name(name)
  raise ConfigurationError, "graph node #{name.inspect} is already declared" if graph_nodes_value.key?(name)
  unless [history, context].all? { |value| value == true || value == false }
    raise ArgumentError, "graph node history and context options must be true or false"
  end
  validate_callable!(input, "node input mapper")

  policies = {timeout:, retries:, retry_on:, retry_delay:}.freeze
  declaration = Node.new(
    name:, assembly:, policies:,
    inherit_history: history,
    inherit_context: context,
    input_mapper: input
  )
  self.graph_nodes_value = graph_nodes_value.merge(name => declaration).freeze
end

# File lib/little_ghost/graph.rb, line 143
def start(name = nil)
  return graph_start_value if name.nil?

  self.graph_start_value = normalize_node_name(name)
end

# File lib/little_ghost/graph.rb, line 281
def to_mermaid
  nodes, edges, error_edges, forks, joins, start_name, finish_name = graph_definition!
  lines = ["flowchart TD"]
  nodes.each_key { |name| lines << "  #{mermaid_id(name)}[#{name}]" }
  lines << "  START((start)) --> #{mermaid_id(start_name)}"
  lines << "  #{mermaid_id(finish_name)} --> FINISH((finish))"
  fan_in_pairs = joins.flat_map { |join| join.from.map { |source| [source, join.to] } }.to_set
  edges.each do |edge|
    next if fan_in_pairs.include?([edge.from, edge.to])

    label = edge.condition ? "condition" : nil
    lines << mermaid_edge(edge.from, edge.to, label:)
  end
  error_edges.each { |edge| lines << mermaid_edge(edge.from, edge.to, label: "error", dotted: true) }
  forks.each do |fork|
    fork.to.each { |target| lines << mermaid_edge(fork.from, target, label: "fork") }
  end
  joins.each do |join|
    join.from.each { |source| lines << mermaid_edge(source, join.to, label: "join") }
  end
  lines.join("\n")
end

# File lib/little_ghost/graph.rb, line 242
def validate!
  graph_definition!
  self
end

# File lib/little_ghost/graph.rb, line 754
def close
  children = @graph_mutex.synchronize { @graph_children.reverse }
  first_error = nil
  children.each do |child|
    child.close
  rescue => error
    first_error ||= error
  end
  super
  raise first_error if first_error
end

# File lib/little_ghost/graph.rb, line 587
def stream(input = nil, history: nil, context: nil,
  cancellation_token: Support::CancellationToken.new, deadline: nil,
  settings: nil, template_locals: nil, template_paths: nil,
  parent_operation_id: nil, checkpoint: nil, **_options)
  raise ArgumentError, "input is required" if input.nil?
  if standalone?
    return build_run(entrypoint_payload(input, {
      history:, context:, settings:, template_paths:,
      deadline_at: deadline, cancellation_token:
    }.compact)).each
  end

  reserve_execution!
  original_input = input.is_a?(Message) ? input : Message.new(role: :user, content: input)
  original_history = normalize_history(history)
  original_context = frozen_state(context || {})
  settings ||= {}
  template_locals ||= {}
  template_paths ||= []

  usage = Usage.new
  error_emitted = false
  Enumerator.new do |events|
    definition = self.class.graph_definition!
    nodes, edges, error_edges, forks, joins, current, finish = definition
    results = {}
    steps = []
    previous = nil
    incoming_edge = nil
    join_context = nil
    routed_error = nil
    previous_step_id = nil

    loop do
      count = next_graph_step!(cancellation_token, deadline)
      state = routing_state(
        input: original_input, history: original_history, context: original_context,
        step: count, current:, previous:, results:,
        predecessors: join_context&.fetch(:predecessors, []),
        incoming_results: join_context&.fetch(:results, {}) || {},
        error: routed_error
      )
      node_input = if join_context
        join_input_for(state, join_context.fetch(:join), nodes.fetch(current))
      else
        node_input_for(state, incoming_edge, nodes.fetch(current))
      end
      terminal = current == finish
      begin
        execution = execute_graph_node(
          node: nodes.fetch(current), input: node_input, history: original_history,
          context: original_context, cancellation_token:, deadline:, settings:,
          template_locals:, template_paths:, parent_operation_id:,
          predecessor_ids: join_context&.fetch(:step_ids, []) || Array(previous_step_id),
          terminal:, events:
        )
      rescue => error
        route = select_error_edge(error_edges.select { |edge| edge.from == current }, error)
        raise unless route

        usage += step_error_usage(error)
        error.instance_variable_set(:@little_ghost_step_usage_accounted, true)
        failed = failed_step(
          error,
          nodes.fetch(current),
          current,
          predecessor_ids: Array(previous_step_id)
        )
        steps << failed
        previous_step_id = failed.id
        incoming_edge = Edge.new(
          from: current, to: route.to, condition: nil,
          input_mapper: route.input_mapper, max_concurrency: nil
        )
        events << transition_event(count, current, route.to, error: true)
        previous = current
        current = route.to
        join_context = nil
        routed_error = error
        next
      end

      results[current] = execution.result
      previous_step_id = execution.step.id
      usage += execution.step.usage
      steps.concat(execution.result.steps)
      if terminal
        final = copy_run_result(execution.result, usage:, steps: steps.freeze)
        execution.events.each do |event|
          event = StreamEvent.build(event.type, **event.data.merge(result: final)) if event.type == :invocation_stop
          error_emitted = true if event.type == :invocation_error
          events << event
        end
        break
      end

      state = routing_state(
        input: original_input, history: original_history, context: original_context,
        step: count, current:, previous:, results:
      )
      selected = select_edge(
        edges.select { |edge| edge.from == current } + forks.select { |fork| fork.from == current },
        state
      )
      if selected.is_a?(Fork)
        fork = selected
        join = joins.find { |candidate| candidate.fork == fork }
        events << StreamEvent.build(
          :assembly_fork,
          assembly_id: self.class.assembly_id,
          assembly_kind: :graph,
          from: current,
          branches: fork.to
        )
        branch_outputs = run_graph_branches(
          fork:, join:, nodes:, edges:, error_edges:,
          original_input:, original_history:, original_context:,
          results:, source_step_id: previous_step_id,
          cancellation_token:, deadline:, settings:,
          template_locals:, template_paths:, parent_operation_id:, events:
        )
        unless join.from.sort == branch_outputs.map(&:terminal).sort
          raise AssemblyRoutingError, "graph fan-out at #{current.inspect} did not reach its inferred fan-in"
        end

        branch_outputs.each do |branch|
          results.merge!(branch.results)
          steps.concat(branch.steps)
          usage += branch.usage
          branch.events.each { |event| events << event }
        end
        events << StreamEvent.build(
          :assembly_join,
          assembly_id: self.class.assembly_id,
          assembly_kind: :graph,
          from: join.from,
          to: join.to
        )
        previous = nil
        current = join.to
        incoming_edge = nil
        join_context = {
          join:,
          predecessors: join.from,
          step_ids: branch_outputs.map { |branch| branch.steps.last.id },
          results: join.from.to_h { |name| [name, results.fetch(name)] }
        }
        previous_step_id = nil
        next
      end

      events << transition_event(count, current, selected.to)
      previous = current
      current = selected.to
      incoming_edge = selected
      join_context = nil
      routed_error = nil
    end
  rescue => error
    usage += unaccounted_step_error_usage(error)
    unless error_emitted
      events << StreamEvent.build(:invocation_error, error:, usage:, metadata: {})
    end
    raise
  end
end