// ============================================================
// Priority and budget scheduler
//
// Scheduler ticks every registered tree. BudgetScheduler is intended for
// larger simulations and agent services where each frame has a hard task
// budget. It chooses the highest-priority runnable tasks, keeps terminal
// tasks dormant until reset, and exposes a report for backpressure metrics.
// Ties are stable because registration order is preserved.
// ============================================================

///|
/// One independently scheduled behavior tree.
pub struct BudgetTask {
  name : String
  mut priority : Int
  root : Node
  bb : Blackboard
  paused : Ref[Bool]
  done : Ref[Bool]
  last : Ref[Status?]
  ticks : Ref[Int]
}

///|
/// Create a scheduled task with an isolated blackboard.
pub fn BudgetTask::new(
  name : String,
  priority : Int,
  root : Node,
  bb : Blackboard,
) -> BudgetTask {
  {
    name,
    priority,
    root,
    bb,
    paused: Ref::new(false),
    done: Ref::new(false),
    last: Ref::new(None),
    ticks: Ref::new(0),
  }
}

///|
/// Task name used in reports and control operations.
pub fn BudgetTask::name(self : BudgetTask) -> String {
  self.name
}

///|
/// Current scheduling priority. Larger values run first.
pub fn BudgetTask::priority(self : BudgetTask) -> Int {
  self.priority
}

///|
/// Change priority without resetting the task.
pub fn BudgetTask::set_priority(self : BudgetTask, priority : Int) -> Unit {
  self.priority = priority
}

///|
/// Return the task's isolated blackboard.
pub fn BudgetTask::blackboard(self : BudgetTask) -> Blackboard {
  self.bb
}

///|
/// Return the last observed status, if the task has run.
pub fn BudgetTask::last_status(self : BudgetTask) -> Status? {
  self.last.get()
}

///|
/// Number of ticks delivered to this task since reset.
pub fn BudgetTask::ticks(self : BudgetTask) -> Int {
  self.ticks.get()
}

///|
/// Whether the task completed and is waiting for an explicit reset.
pub fn BudgetTask::is_done(self : BudgetTask) -> Bool {
  self.done.get()
}

///|
/// Whether the task is paused.
pub fn BudgetTask::is_paused(self : BudgetTask) -> Bool {
  self.paused.get()
}

///|
/// One task result from a budget scheduler frame.
pub struct BudgetTick {
  name : String
  priority : Int
  frame : Int
  status : Status
}

///|
/// Name of the task that was ticked.
pub fn BudgetTick::name(self : BudgetTick) -> String {
  self.name
}

///|
/// Priority used for this result.
pub fn BudgetTick::priority(self : BudgetTick) -> Int {
  self.priority
}

///|
/// Scheduler frame in which the result was produced.
pub fn BudgetTick::frame(self : BudgetTick) -> Int {
  self.frame
}

///|
/// Status returned by the task.
pub fn BudgetTick::status(self : BudgetTick) -> Status {
  self.status
}

///|
/// Stable CSV row for a task result.
pub fn BudgetTick::to_csv(self : BudgetTick) -> String {
  "\{self.frame},\{self.name},\{self.priority},\{self.status.to_string()}"
}

///|
/// Aggregate outcome of one budget scheduler frame.
pub struct BudgetReport {
  frame : Int
  attempted : Int
  skipped : Int
  pending : Int
  results : Array[BudgetTick]
}

///|
/// Current scheduler frame.
pub fn BudgetReport::frame(self : BudgetReport) -> Int {
  self.frame
}

///|
/// Number of tasks ticked in this frame.
pub fn BudgetReport::attempted(self : BudgetReport) -> Int {
  self.attempted
}

///|
/// Number of runnable tasks left out by the budget.
pub fn BudgetReport::skipped(self : BudgetReport) -> Int {
  self.skipped
}

///|
/// Number of tasks not yet terminal, including paused tasks.
pub fn BudgetReport::pending(self : BudgetReport) -> Int {
  self.pending
}

///|
/// Copy the results in scheduling order.
pub fn BudgetReport::results(self : BudgetReport) -> Array[BudgetTick] {
  let copy : Array[BudgetTick] = []
  for result in self.results {
    copy.push(result)
  }
  copy
}

///|
/// Export the frame report as CSV.
pub fn BudgetReport::to_csv(self : BudgetReport) -> String {
  let output = StringBuilder::new()
  output.write_string("frame,name,priority,status\n")
  for result in self.results {
    output.write_string(result.to_csv())
    output.write_string("\n")
  }
  output.to_string()
}

///|
/// A scheduler with a hard per-frame task budget.
pub struct BudgetScheduler {
  tasks : Array[BudgetTask]
  budget : Ref[Int]
  frame : Ref[Int]
}

///|
/// Create an empty priority scheduler.
pub fn BudgetScheduler::new(max_tasks_per_frame : Int) -> BudgetScheduler {
  {
    tasks: [],
    budget: Ref::new(
      if max_tasks_per_frame > 0 {
        max_tasks_per_frame
      } else {
        1
      },
    ),
    frame: Ref::new(0),
  }
}

///|
/// Return the maximum number of tasks per frame.
pub fn BudgetScheduler::budget(self : BudgetScheduler) -> Int {
  self.budget.get()
}

///|
/// Update the frame budget; invalid values become one.
pub fn BudgetScheduler::set_budget(
  self : BudgetScheduler,
  max_tasks_per_frame : Int,
) -> Unit {
  self.budget.set(if max_tasks_per_frame > 0 { max_tasks_per_frame } else { 1 })
}

///|
/// Add a task unless another task has the same name.
pub fn BudgetScheduler::add(self : BudgetScheduler, task : BudgetTask) -> Bool {
  if self.has(task.name) {
    false
  } else {
    self.tasks.push(task)
    true
  }
}

///|
/// Remove a named task. The task's node is not reset after removal.
pub fn BudgetScheduler::remove(self : BudgetScheduler, name : String) -> Bool {
  let mut i = 0
  while i < self.tasks.length() {
    if self.tasks[i].name == name {
      let _ = self.tasks.remove(i)
      return true
    }
    i = i + 1
  }
  false
}

///|
/// Return whether a named task exists.
pub fn BudgetScheduler::has(self : BudgetScheduler, name : String) -> Bool {
  for task in self.tasks {
    if task.name == name {
      return true
    }
  }
  false
}

///|
/// Number of registered tasks.
pub fn BudgetScheduler::size(self : BudgetScheduler) -> Int {
  self.tasks.length()
}

///|
/// Pause one task without changing its node state.
pub fn BudgetScheduler::pause(self : BudgetScheduler, name : String) -> Bool {
  for task in self.tasks {
    if task.name == name {
      task.paused.set(true)
      return true
    }
  }
  false
}

///|
/// Resume one paused task.
pub fn BudgetScheduler::resume_task(
  self : BudgetScheduler,
  name : String,
) -> Bool {
  for task in self.tasks {
    if task.name == name {
      task.paused.set(false)
      return true
    }
  }
  false
}

///|
/// Change one task's priority.
pub fn BudgetScheduler::set_priority(
  self : BudgetScheduler,
  name : String,
  priority : Int,
) -> Bool {
  for task in self.tasks {
    if task.name == name {
      task.set_priority(priority)
      return true
    }
  }
  false
}

///|
/// Reset one task and make it runnable again.
pub fn BudgetScheduler::reset(self : BudgetScheduler, name : String) -> Bool {
  for task in self.tasks {
    if task.name == name {
      task.root.reset()
      task.done.set(false)
      task.last.set(None)
      task.ticks.set(0)
      return true
    }
  }
  false
}

///|
/// Reset all tasks and the scheduler frame counter.
pub fn BudgetScheduler::reset_all(self : BudgetScheduler) -> Unit {
  for task in self.tasks {
    task.root.reset()
    task.done.set(false)
    task.last.set(None)
    task.ticks.set(0)
  }
  self.frame.set(0)
}

///|
/// Number of tasks that have not reached a terminal result.
pub fn BudgetScheduler::pending(self : BudgetScheduler) -> Int {
  let mut count = 0
  for task in self.tasks {
    if !task.done.get() {
      count = count + 1
    }
  }
  count
}

///|
/// Choose the highest-priority unselected runnable task.
fn choose_budget_task(tasks : Array[BudgetTask], selected : Array[Bool]) -> Int {
  let mut best = -1
  let mut best_priority = -2147483647
  let mut i = 0
  while i < tasks.length() {
    let task = tasks[i]
    if !selected[i] && !task.paused.get() && !task.done.get() {
      if best < 0 || task.priority > best_priority {
        best = i
        best_priority = task.priority
      }
    }
    i = i + 1
  }
  best
}

///|
/// Tick the highest-priority runnable tasks up to the frame budget.
pub fn BudgetScheduler::tick(self : BudgetScheduler) -> BudgetReport {
  self.frame.set(self.frame.get() + 1)
  let selected : Array[Bool] = []
  for _ in self.tasks {
    selected.push(false)
  }
  let mut runnable_before = 0
  for task in self.tasks {
    if !task.paused.get() && !task.done.get() {
      runnable_before = runnable_before + 1
    }
  }
  let results : Array[BudgetTick] = []
  let mut attempted = 0
  while attempted < self.budget.get() {
    let index = choose_budget_task(self.tasks, selected)
    if index < 0 {
      break
    }
    selected[index] = true
    let task = self.tasks[index]
    let status = task.root.tick(task.bb)
    task.last.set(Some(status))
    task.ticks.set(task.ticks.get() + 1)
    match status {
      Status::BTSuccess | Status::BTFailure => task.done.set(true)
      Status::BTRunning => ()
    }
    results.push({
      name: task.name,
      priority: task.priority,
      frame: self.frame.get(),
      status,
    })
    attempted = attempted + 1
  }
  let pending = self.pending()
  {
    frame: self.frame.get(),
    attempted,
    skipped: if runnable_before > attempted {
      runnable_before - attempted
    } else {
      0
    },
    pending,
    results,
  }
}

///|
/// Current scheduler frame number.
pub fn BudgetScheduler::frame(self : BudgetScheduler) -> Int {
  self.frame.get()
}

///|
/// Return a copy of all registered task names.
pub fn BudgetScheduler::task_names(self : BudgetScheduler) -> Array[String] {
  let names : Array[String] = []
  for task in self.tasks {
    names.push(task.name)
  }
  names
}

///|
/// Return the first task with the requested name, if any.
pub fn BudgetScheduler::task(
  self : BudgetScheduler,
  name : String,
) -> BudgetTask? {
  for task in self.tasks {
    if task.name == name {
      return Some(task)
    }
  }
  None
}

///|
/// Remove all tasks and reset the frame counter.
pub fn BudgetScheduler::clear(self : BudgetScheduler) -> Unit {
  self.tasks.clear()
  self.frame.set(0)
}