// TinyRTOS — a tiny C++ priority-preemptive kernel you can actually compile.
//
//   g++ -std=c++17 -Wall kernel.cpp -o tinyr-tos && ./tinyr-tos
//   ./tinyr-tos invert     # priority inversion (the bug)
//   ./tinyr-tos inherit    # same bug, with priority inheritance (the fix)
//
// What this is
// ------------
// A tick-accurate model of a fixed-priority preemptive scheduler.
// Every tick, pick() runs the highest-priority READY task. That is the
// entire policy. FreeRTOS does this. Linux SCHED_FIFO does this.
// A Network Element Controller does this so a Loss-of-Signal alarm can
// steal the CPU from telemetry the instant the fiber goes dark.
//
// Mapping onto real Linux (PREEMPT_RT / embedded):
//   pthread_setschedparam(t, SCHED_FIFO, { .sched_priority = 99 - prio });
//   A hardware interrupt notifies a task (see Kernel::notify); it does
//   not do the work itself — classic deferred interrupt processing.
//
// Priority: 0 = highest  (same convention as this playground)

#include <cstdio>
#include <string>

enum class State { Ready, Running, Blocked };

struct Task {
  const char* name;
  int prio;            // base priority, 0 = highest
  int effective;       // may be boosted (priority inheritance)
  State state;
  int work_left;       // remaining CPU ticks in this job
  bool infinite;       // background tasks never go idle
  int wait_mutex;      // mutex id, or -1
  int wait_event;      // 1 = sleeping on an event (LOS)
  int deadline;        // tick the current job must finish by, or -1
};

struct Mutex {
  int holder;          // task index, or -1
  int waiting[8];
  int nwait;
};

class Kernel {
 public:
  static constexpr int kMaxTasks = 8;

  Task tasks[kMaxTasks];
  int ntasks = 0;
  int current = -1;
  int tick = 0;
  bool inherit = false;
  Mutex mutexes[2];

  Kernel() {
    mutexes[0].holder = mutexes[1].holder = -1;
    mutexes[0].nwait = mutexes[1].nwait = 0;
  }

  int spawn(const char* name, int prio, int work, bool infinite = false) {
    Task& t = tasks[ntasks];
    t.name = name;
    t.prio = t.effective = prio;
    t.state = State::Ready;
    t.work_left = work;
    t.infinite = infinite;
    t.wait_mutex = -1;
    t.wait_event = 0;
    t.deadline = -1;
    return ntasks++;
  }

  bool runnable(int i) const {
    const Task& t = tasks[i];
    if (t.state != State::Ready && t.state != State::Running) return false;
    return t.infinite || t.work_left > 0;
  }

  // Always run the highest-priority ready task. That's the RTOS.
  int pick() {
    int best = -1;
    for (int i = 0; i < ntasks; i++) {
      if (!runnable(i)) continue;
      if (best < 0 || tasks[i].effective < tasks[best].effective)
        best = i;
      else if (tasks[i].effective == tasks[best].effective && i == current)
        best = i;  // stable: don't bounce between equal priorities
    }
    return best;
  }

  // ISR / another task wakes a sleeper. The sleeper does NOT run here —
  // it just becomes READY. pick() will choose it on the next schedule.
  void notify(int tid) {
    Task& t = tasks[tid];
    if (t.wait_event) {
      t.wait_event = 0;
      t.state = State::Ready;
    }
  }

  void lock(int tid, int mid) {
    Mutex& m = mutexes[mid];
    if (m.holder < 0) {
      m.holder = tid;
      return;
    }
    Task& waiter = tasks[tid];
    waiter.state = State::Blocked;
    waiter.wait_mutex = mid;
    m.waiting[m.nwait++] = tid;
    if (inherit) {
      int& holder_eff = tasks[m.holder].effective;
      if (waiter.effective < holder_eff) holder_eff = waiter.effective;
    }
  }

  void unlock(int tid, int mid) {
    Mutex& m = mutexes[mid];
    if (m.holder != tid) return;
    m.holder = -1;
    tasks[tid].effective = tasks[tid].prio;  // drop boost
    if (m.nwait == 0) return;

    int best_i = 0;
    for (int i = 1; i < m.nwait; i++)
      if (tasks[m.waiting[i]].prio < tasks[m.waiting[best_i]].prio) best_i = i;
    int next = m.waiting[best_i];
    for (int i = best_i; i < m.nwait - 1; i++) m.waiting[i] = m.waiting[i + 1];
    m.nwait--;
    m.holder = next;
    tasks[next].wait_mutex = -1;
    tasks[next].state = State::Ready;
  }

  // One timer tick. Returns whoever ran, or "IDLE" / "ISR".
  const char* step(bool isr = false) {
    if (isr) {
      tick++;
      return "ISR";
    }
    int next = pick();
    if (current >= 0 && current != next && tasks[current].state == State::Running)
      tasks[current].state = State::Ready;  // preempted
    current = next;
    tick++;
    if (current < 0) return "IDLE";
    tasks[current].state = State::Running;
    if (tasks[current].work_left > 0) tasks[current].work_left--;
    if (!tasks[current].infinite && tasks[current].work_left == 0)
      tasks[current].state = State::Blocked;
    return tasks[current].name;
  }
};

static void banner(const char* title) {
  std::printf("\n== %s ==\n", title);
  std::printf("tick  CPU            note\n");
  std::printf("----  -------------  ----\n");
}

// ---------------------------------------------------------------------------
// Scenario 1: fiber cut on an optical Network Element.
// Framer/Telemetry/CLI are always READY. LOS sleeps until the interrupt.
// ---------------------------------------------------------------------------
static void run_fiber() {
  Kernel k;
  k.spawn("FRAMER", 1, 0, true);
  k.spawn("TELEMETRY", 2, 0, true);
  k.spawn("CLI", 3, 0, true);
  int los = k.spawn("LOS", 0, 0, false);

  k.tasks[los].state = State::Blocked;
  k.tasks[los].wait_event = 1;

  const int cut_at = 8;
  const int los_work = 3;
  const int los_slack = 4;  // must finish within 4 ticks of the ISR

  banner("fiber cut  (priority-preemptive)");
  bool met = false;
  for (int t = 0; t < 20; t++) {
    const char* who;
    const char* note = "";
    if (t == cut_at) {
      who = k.step(true);  // hardware ISR — not a task
      k.tasks[los].work_left = los_work;
      k.tasks[los].deadline = k.tick + los_slack;
      k.notify(los);
      note = "LOS interrupt (fiber cut) — notify LOS task";
    } else {
      who = k.step();
      if (std::string(who) == "LOS" && k.tasks[los].work_left == 0 && !met) {
        met = true;
        note = k.tick <= k.tasks[los].deadline ? "deadline MET" : "deadline MISS";
        k.tasks[los].wait_event = 1;
        k.tasks[los].deadline = -1;
      }
    }
    std::printf("%4d  %-13s  %s\n", t, who, note);
  }
}

// ---------------------------------------------------------------------------
// Scenario 2: priority inversion.
// LOW holds a mutex. HIGH needs it. MED (who doesn't) runs instead of HIGH.
// Inheritance boosts LOW up to HIGH's priority so MED cannot sneak in.
//
// LOW:  run, lock, run x4, unlock
// HIGH: (arrives t=3) run, lock, run x2
// MED:  (arrives t=4) run x6
// ---------------------------------------------------------------------------
static void run_inversion(bool inherit) {
  Kernel k;
  k.inherit = inherit;
  int high = k.spawn("HIGH", 0, 3);
  int med  = k.spawn("MED",  1, 6);
  int low  = k.spawn("LOW",  2, 6);

  k.tasks[high].state = State::Blocked;
  k.tasks[med].state  = State::Blocked;

  banner(inherit ? "inversion + inheritance (the fix)"
                 : "priority inversion (the bug)");

  int low_ran = 0, high_ran = 0;
  bool low_has_lock = false, high_tried = false;

  for (int t = 0; t < 20; t++) {
    if (t == 3) k.tasks[high].state = State::Ready;
    if (t == 4) k.tasks[med].state  = State::Ready;

    const char* note = "";

    const char* who = k.step();

    if (std::string(who) == "LOW") {
      low_ran++;
      if (low_ran == 1) {
        k.lock(low, 0);
        low_has_lock = true;
        note = "LOW locks mutex";
      } else if (low_has_lock && low_ran == 5) {
        k.unlock(low, 0);
        low_has_lock = false;
        note = "LOW unlocks — HIGH can run";
      }
    }
    if (std::string(who) == "HIGH") {
      high_ran++;
      if (high_ran == 1 && !high_tried) {
        high_tried = true;
        k.lock(high, 0);
        if (k.tasks[high].state == State::Blocked)
          note = inherit ? "HIGH blocks; LOW inherits prio 0"
                         : "HIGH blocks on mutex — INVERSION starts";
      }
    }
    std::printf("%4d  %-13s  %s\n", t, who, note);
  }
}

int main(int argc, char** argv) {
  const char* mode = argc > 1 ? argv[1] : "fiber";
  std::printf("TinyRTOS — priority-preemptive C++ kernel  (prio 0 = highest)\n");
  if (std::string(mode) == "invert") run_inversion(false);
  else if (std::string(mode) == "inherit") run_inversion(true);
  else run_fiber();
  return 0;
}