///|
/// Task identifier used by the public graph API.
pub struct TaskId {
  priv value : String
} derive(Eq, Debug)

///|
/// Execution status tracked for each task node.
pub(all) enum TaskStatus {
  Pending
  Ready
  Running
  Succeeded
  Failed(String)
  Skipped(String)
} derive(Eq, Debug)

///|
/// A node in a reproducible research or agent workflow.
pub struct TaskNode {
  priv id : TaskId
  priv title : String
  priv description : String
  priv inputs : Array[String]
  priv outputs : Array[String]
  priv tags : Array[String]
  priv status : TaskStatus
} derive(Eq, Debug)

///|
/// Directed dependency edge: `before` must finish before `after`.
pub struct Dependency {
  priv before : TaskId
  priv after : TaskId
} derive(Eq, Debug)

///|
/// A validated execution plan with serial order and parallel-ready batches.
pub struct ExecutionPlan {
  priv order : Array[TaskId]
  priv batches : Array[Array[TaskId]]
} derive(Eq, Debug)

///|
/// Graph errors retain the ids and edges needed for useful diagnostics.
pub(all) enum GraphError {
  DuplicateTask(TaskId)
  DuplicateDependency(Dependency)
  MissingTask(TaskId)
  MissingDependencyEndpoint(Dependency)
  CycleDetected(Array[TaskId])
} derive(Eq, Debug)

///|
/// Errors raised by checked task-status transitions.
pub(all) enum StatusTransitionError {
  TransitionMissingTask(TaskId)
  InvalidStatusTransition(TaskId, TaskStatus, TaskStatus)
} derive(Eq, Debug)

///|
/// Errors raised when a plan or trace no longer matches its graph.
pub(all) enum SnapshotError {
  SnapshotGraphError(GraphError)
  StaleExecutionPlan
  UnknownTraceTask(TaskId)
} derive(Eq, Debug)

///|
/// Mutable workflow graph.
pub struct FlowGraph {
  priv mut tasks : Array[TaskNode]
  priv mut dependencies : Array[Dependency]
} derive(Debug)

///|
/// Build a task id.
pub fn TaskId::new(value : String) -> TaskId {
  { value, }
}

///|
/// Return the raw string value of this task id.
pub fn TaskId::value(self : TaskId) -> String {
  self.value
}

///|
/// Build a basic pending task.
pub fn TaskNode::new(id : String, title : String) -> TaskNode {
  {
    id: TaskId::new(id),
    title,
    description: "",
    inputs: [],
    outputs: [],
    tags: [],
    status: Pending,
  }
}

///|
/// Return this task's id.
pub fn TaskNode::id(self : TaskNode) -> TaskId {
  self.id
}

///|
/// Return this task's title.
pub fn TaskNode::title(self : TaskNode) -> String {
  self.title
}

///|
/// Return this task's description.
pub fn TaskNode::description(self : TaskNode) -> String {
  self.description
}

///|
/// Return a detached copy of this task's input labels.
pub fn TaskNode::inputs(self : TaskNode) -> Array[String] {
  self.inputs.copy()
}

///|
/// Return a detached copy of this task's output labels.
pub fn TaskNode::outputs(self : TaskNode) -> Array[String] {
  self.outputs.copy()
}

///|
/// Return a detached copy of this task's tags.
pub fn TaskNode::tags(self : TaskNode) -> Array[String] {
  self.tags.copy()
}

///|
/// Return this task's execution status.
pub fn TaskNode::status(self : TaskNode) -> TaskStatus {
  self.status
}

///|
/// Add a short description.
pub fn TaskNode::with_description(
  self : TaskNode,
  description : String,
) -> TaskNode {
  { ..self, description, }
}

///|
/// Add input labels or artifacts.
pub fn TaskNode::with_inputs(
  self : TaskNode,
  inputs : Array[String],
) -> TaskNode {
  { ..self, inputs, }
}

///|
/// Add output labels or artifacts.
pub fn TaskNode::with_outputs(
  self : TaskNode,
  outputs : Array[String],
) -> TaskNode {
  { ..self, outputs, }
}

///|
/// Add topic tags.
pub fn TaskNode::with_tags(self : TaskNode, tags : Array[String]) -> TaskNode {
  { ..self, tags, }
}

///|
/// Override status.
pub fn TaskNode::with_status(self : TaskNode, status : TaskStatus) -> TaskNode {
  { ..self, status, }
}

///|
/// Build a dependency edge.
pub fn Dependency::new(before : TaskId, after : TaskId) -> Dependency {
  { before, after }
}

///|
/// Return the predecessor endpoint of this dependency.
pub fn Dependency::before(self : Dependency) -> TaskId {
  self.before
}

///|
/// Return the successor endpoint of this dependency.
pub fn Dependency::after(self : Dependency) -> TaskId {
  self.after
}

///|
/// Return a detached copy of the serial execution order.
pub fn ExecutionPlan::order(self : ExecutionPlan) -> Array[TaskId] {
  self.order.copy()
}

///|
/// Return a detached copy of the parallel execution batches.
pub fn ExecutionPlan::batches(self : ExecutionPlan) -> Array[Array[TaskId]] {
  let batches : Array[Array[TaskId]] = []
  for batch in self.batches {
    batches.push(batch.copy())
  }
  batches
}

///|
/// Create an empty flow graph.
pub fn FlowGraph::new() -> FlowGraph {
  { tasks: [], dependencies: [] }
}

///|
/// Return a detached copy of task nodes and their metadata arrays.
pub fn FlowGraph::tasks(self : FlowGraph) -> Array[TaskNode] {
  self.tasks_snapshot()
}

///|
/// Return a detached copy of dependency edges.
pub fn FlowGraph::dependencies(self : FlowGraph) -> Array[Dependency] {
  self.dependencies_snapshot()
}

///|
/// Return a detached copy of task nodes and their metadata arrays.
pub fn FlowGraph::tasks_snapshot(self : FlowGraph) -> Array[TaskNode] {
  let out : Array[TaskNode] = []
  for task in self.tasks {
    out.push(task.snapshot())
  }
  out
}

///|
/// Return a detached copy of dependency edges.
pub fn FlowGraph::dependencies_snapshot(self : FlowGraph) -> Array[Dependency] {
  self.dependencies.copy()
}

///|
/// Return a detached copy of this graph and all task metadata arrays.
pub fn FlowGraph::snapshot(self : FlowGraph) -> FlowGraph {
  let tasks : Array[TaskNode] = []
  for task in self.tasks {
    tasks.push(task.snapshot())
  }
  { tasks, dependencies: self.dependencies.copy() }
}

///|
/// Return a detached copy of a task node and its metadata arrays.
pub fn TaskNode::snapshot(self : TaskNode) -> TaskNode {
  {
    ..self,
    inputs: self.inputs.copy(),
    outputs: self.outputs.copy(),
    tags: self.tags.copy(),
  }
}

///|
/// Return a detached copy of an execution plan, including nested batches.
pub fn ExecutionPlan::snapshot(self : ExecutionPlan) -> ExecutionPlan {
  let batches : Array[Array[TaskId]] = []
  for batch in self.batches {
    batches.push(batch.copy())
  }
  { order: self.order.copy(), batches }
}

///|
/// Return the number of task nodes.
pub fn FlowGraph::task_count(self : FlowGraph) -> Int {
  self.tasks.length()
}

///|
/// Return the number of dependency edges.
pub fn FlowGraph::dependency_count(self : FlowGraph) -> Int {
  self.dependencies.length()
}

///|
/// Return tasks with no incoming dependency edges.
pub fn FlowGraph::roots(self : FlowGraph) -> Result[Array[TaskId], GraphError] {
  match self.validate() {
    Err(err) => Err(err)
    Ok(_) => {
      let roots : Array[TaskId] = []
      for task in self.tasks {
        let mut has_predecessor = false
        for dep in self.dependencies {
          if dep.after == task.id {
            has_predecessor = true
          }
        }
        if !has_predecessor {
          roots.push(task.id)
        }
      }
      Ok(roots)
    }
  }
}

///|
/// Return tasks with no outgoing dependency edges.
pub fn FlowGraph::leaves(self : FlowGraph) -> Result[Array[TaskId], GraphError] {
  match self.validate() {
    Err(err) => Err(err)
    Ok(_) => {
      let leaves : Array[TaskId] = []
      for task in self.tasks {
        if self.successor_ids(task.id).is_empty() {
          leaves.push(task.id)
        }
      }
      Ok(leaves)
    }
  }
}

///|
/// Add a task node. Duplicate ids return `DuplicateTask`.
pub fn FlowGraph::add_task(
  self : FlowGraph,
  node : TaskNode,
) -> Result[Unit, GraphError] {
  if self.contains_task(node.id) {
    return Err(DuplicateTask(node.id))
  }
  self.tasks.push(node.snapshot())
  Ok(())
}

///|
/// Add a dependency edge. Duplicate edges are rejected; endpoint existence is
/// checked by `validate`.
pub fn FlowGraph::add_dependency(
  self : FlowGraph,
  before : TaskId,
  after : TaskId,
) -> Result[Unit, GraphError] {
  let dep = { before, after }
  for current in self.dependencies {
    if current == dep {
      return Err(DuplicateDependency(dep))
    }
  }
  self.dependencies.push(dep)
  Ok(())
}

///|
/// Update the status of an existing task.
pub fn FlowGraph::update_status(
  self : FlowGraph,
  id : TaskId,
  status : TaskStatus,
) -> Result[Unit, GraphError] {
  for i = 0; i < self.tasks.length(); i = i + 1 {
    if self.tasks[i].id == id {
      self.tasks[i] = { ..self.tasks[i], status, }
      return Ok(())
    }
  }
  Err(MissingTask(id))
}

///|
/// Update a task through the documented execution-state machine.
///
/// `update_status` remains available for replay and compatibility. Normal
/// execution should prefer this checked method.
pub fn FlowGraph::transition_status(
  self : FlowGraph,
  id : TaskId,
  status : TaskStatus,
) -> Result[Unit, StatusTransitionError] {
  for i = 0; i < self.tasks.length(); i = i + 1 {
    if self.tasks[i].id == id {
      let current = self.tasks[i].status
      if !current.can_transition_to(status) {
        return Err(InvalidStatusTransition(id, current, status))
      }
      self.tasks[i] = { ..self.tasks[i], status, }
      return Ok(())
    }
  }
  Err(TransitionMissingTask(id))
}

///|
/// Return whether a status transition is valid for normal execution.
pub fn TaskStatus::can_transition_to(
  self : TaskStatus,
  next : TaskStatus,
) -> Bool {
  if self == next {
    return true
  }
  match (self, next) {
    (Pending, Ready) | (Pending, Running) | (Pending, Skipped(_)) => true
    (Ready, Running) | (Ready, Skipped(_)) => true
    (Running, Succeeded) | (Running, Failed(_)) => true
    (Failed(_), Ready) => true
    _ => false
  }
}

///|
/// Return immediate predecessor ids for a task.
pub fn FlowGraph::predecessors(
  self : FlowGraph,
  id : TaskId,
) -> Result[Array[TaskId], GraphError] {
  if !self.contains_task(id) {
    return Err(MissingTask(id))
  }
  let out : Array[TaskId] = []
  for dep in self.dependencies {
    if dep.after == id {
      out.push(dep.before)
    }
  }
  Ok(out)
}

///|
/// Return immediate successor ids for a task.
pub fn FlowGraph::successors(
  self : FlowGraph,
  id : TaskId,
) -> Result[Array[TaskId], GraphError] {
  if !self.contains_task(id) {
    return Err(MissingTask(id))
  }
  Ok(self.successor_ids(id))
}

///|
/// Return tasks whose predecessors are all in `done`.
pub fn FlowGraph::ready_tasks(
  self : FlowGraph,
  done : Array[TaskId],
) -> Result[Array[TaskId], GraphError] {
  match self.validate() {
    Err(err) => Err(err)
    Ok(_) => {
      for id in done {
        if !self.contains_task(id) {
          return Err(MissingTask(id))
        }
      }
      let ready : Array[TaskId] = []
      for task in self.tasks {
        if !contains_id(done, task.id) &&
          self.all_predecessors_consumed(task.id, done) {
          ready.push(task.id)
        }
      }
      Ok(ready)
    }
  }
}

///|
/// Return pending or ready tasks whose predecessors have succeeded.
///
/// Failed or skipped predecessors block their successors. Callers that manage
/// completion separately should continue using `ready_tasks(done)`.
pub fn FlowGraph::runnable_tasks(
  self : FlowGraph,
) -> Result[Array[TaskId], GraphError] {
  match self.validate() {
    Err(err) => Err(err)
    Ok(_) => {
      let runnable : Array[TaskId] = []
      for task in self.tasks {
        if task.status.is_runnable() && self.all_predecessors_succeeded(task.id) {
          runnable.push(task.id)
        }
      }
      Ok(runnable)
    }
  }
}

///|
/// Validate all edge endpoints and cycle freedom.
pub fn FlowGraph::validate(self : FlowGraph) -> Result[Unit, GraphError] {
  for dep in self.dependencies {
    if !self.contains_task(dep.before) || !self.contains_task(dep.after) {
      return Err(MissingDependencyEndpoint(dep))
    }
  }
  match self.find_cycle() {
    Some(path) => Err(CycleDetected(path))
    None => Ok(())
  }
}

///|
/// Return a topological execution order.
pub fn FlowGraph::topological_sort(
  self : FlowGraph,
) -> Result[Array[TaskId], GraphError] {
  match self.validate() {
    Err(err) => Err(err)
    Ok(_) => {
      let consumed : Array[TaskId] = []
      let order : Array[TaskId] = []
      while order.length() < self.tasks.length() {
        let mut progressed = false
        for task in self.tasks {
          if !contains_id(order, task.id) &&
            self.all_predecessors_consumed(task.id, consumed) {
            order.push(task.id)
            consumed.push(task.id)
            progressed = true
          }
        }
        if !progressed {
          return Err(CycleDetected(order))
        }
      }
      Ok(order)
    }
  }
}

///|
/// Return parallel-ready batches. Tasks in the same batch have no dependencies between them.
pub fn FlowGraph::execution_batches(
  self : FlowGraph,
) -> Result[Array[Array[TaskId]], GraphError] {
  match self.validate() {
    Err(err) => Err(err)
    Ok(_) => {
      let done : Array[TaskId] = []
      let batches : Array[Array[TaskId]] = []
      while done.length() < self.tasks.length() {
        let batch : Array[TaskId] = []
        for task in self.tasks {
          if !contains_id(done, task.id) &&
            self.all_predecessors_consumed(task.id, done) {
            batch.push(task.id)
          }
        }
        if batch.is_empty() {
          return Err(CycleDetected(done))
        }
        for id in batch {
          done.push(id)
        }
        batches.push(batch)
      }
      Ok(batches)
    }
  }
}

///|
/// Create a full execution plan.
pub fn FlowGraph::plan(self : FlowGraph) -> Result[ExecutionPlan, GraphError] {
  match self.topological_sort() {
    Err(err) => Err(err)
    Ok(order) =>
      match self.execution_batches() {
        Err(err) => Err(err)
        Ok(batches) => Ok({ order, batches })
      }
  }
}

///|
/// Check that a plan still exactly matches the graph's deterministic plan.
pub fn FlowGraph::validate_plan(
  self : FlowGraph,
  plan : ExecutionPlan,
) -> Result[Unit, SnapshotError] {
  match self.plan() {
    Err(err) => Err(SnapshotGraphError(err))
    Ok(expected) =>
      if expected == plan {
        Ok(())
      } else {
        Err(StaleExecutionPlan)
      }
  }
}

///|
/// Check that every trace event refers to a task in this graph.
pub fn FlowGraph::validate_trace(
  self : FlowGraph,
  trace : Trace,
) -> Result[Unit, SnapshotError] {
  for event in trace.events {
    if !self.contains_task(event.task_id) {
      return Err(UnknownTraceTask(event.task_id))
    }
  }
  Ok(())
}

///|
/// Check graph validity, plan freshness, and trace task references together.
pub fn FlowGraph::validate_snapshot(
  self : FlowGraph,
  plan : ExecutionPlan,
  trace : Trace,
) -> Result[Unit, SnapshotError] {
  match self.validate_plan(plan) {
    Err(err) => Err(err)
    Ok(_) => self.validate_trace(trace)
  }
}

///|
/// Find a task by id.
pub fn FlowGraph::find_task(self : FlowGraph, id : TaskId) -> TaskNode? {
  for task in self.tasks {
    if task.id == id {
      return Some(task.snapshot())
    }
  }
  None
}

///|
fn FlowGraph::contains_task(self : FlowGraph, id : TaskId) -> Bool {
  for task in self.tasks {
    if task.id == id {
      return true
    }
  }
  false
}

///|
fn FlowGraph::all_predecessors_consumed(
  self : FlowGraph,
  id : TaskId,
  consumed : Array[TaskId],
) -> Bool {
  for dep in self.dependencies {
    if dep.after == id && !contains_id(consumed, dep.before) {
      return false
    }
  }
  true
}

///|
fn FlowGraph::all_predecessors_succeeded(self : FlowGraph, id : TaskId) -> Bool {
  for dep in self.dependencies {
    if dep.after == id {
      match self.find_task(dep.before) {
        Some(task) => if task.status != Succeeded { return false }
        None => return false
      }
    }
  }
  true
}

///|
fn TaskStatus::is_runnable(self : TaskStatus) -> Bool {
  match self {
    Pending | Ready => true
    _ => false
  }
}

///|
fn FlowGraph::successor_ids(self : FlowGraph, id : TaskId) -> Array[TaskId] {
  let out : Array[TaskId] = []
  for dep in self.dependencies {
    if dep.before == id {
      out.push(dep.after)
    }
  }
  out
}

///|
fn FlowGraph::find_cycle(self : FlowGraph) -> Array[TaskId]? {
  let visiting : Array[TaskId] = []
  let visited : Array[TaskId] = []
  for task in self.tasks {
    match self.visit_for_cycle(task.id, visiting, visited) {
      Some(path) => return Some(path)
      None => continue
    }
  }
  None
}

///|
fn FlowGraph::visit_for_cycle(
  self : FlowGraph,
  id : TaskId,
  visiting : Array[TaskId],
  visited : Array[TaskId],
) -> Array[TaskId]? {
  if contains_id(visited, id) {
    return None
  }
  if contains_id(visiting, id) {
    return Some(cycle_suffix(visiting, id))
  }
  visiting.push(id)
  for next in self.successor_ids(id) {
    match self.visit_for_cycle(next, visiting, visited) {
      Some(path) => return Some(path)
      None => continue
    }
  }
  let _ = visiting.pop()
  visited.push(id)
  None
}

///|
fn contains_id(ids : Array[TaskId], id : TaskId) -> Bool {
  for current in ids {
    if current == id {
      return true
    }
  }
  false
}

///|
fn cycle_suffix(ids : Array[TaskId], repeated : TaskId) -> Array[TaskId] {
  let out : Array[TaskId] = []
  let mut found = false
  for id in ids {
    if id == repeated {
      found = true
    }
    if found {
      out.push(id)
    }
  }
  out.push(repeated)
  out
}