#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SH110X.h>

// ===== Display setup =====
#define I2C_SDA 21
#define I2C_SCL 22

Adafruit_SH1107 display(128, 128, &Wire);

// ===== VM OPCODES =====
enum {
  OP_PUSH = 1,
  OP_ADD,
  OP_SUB,
  OP_MUL,
  OP_DIV,
  OP_PRINT,   // legacy, direct print
  OP_LOAD,
  OP_STORE,
  OP_JMP,
  OP_JZ,
  OP_SYS,     // syscall
  OP_HALT = 99
};

// ===== Syscall numbers =====
enum {
  SYS_PRINT_INT = 1
};

// ===== VM / Process definition =====
struct Process {
  const int *code;
  int codeLen;

  int ip;          // instruction pointer
  int stack[64];
  int sp;

  int vars[16];

  bool alive;
  int pid;
};

const int MAX_PROCS = 4;
Process procs[MAX_PROCS];
int procCount = 0;
int currentProc = 0;
int nextPid = 1;

// ===== System state =====
int cursorY = 0;
bool firmwareCrashed = false;
bool autoScheduling = true;  // can be paused via serial

// ===== Simple OLED text output =====
void oledClear() {
  display.clearDisplay();
  cursorY = 0;
  display.display();
}

void oledPrintLine(const String &s) {
  if (cursorY > 120) {
    display.clearDisplay();
    cursorY = 0;
  }
  display.setCursor(0, cursorY);
  display.print(s);
  display.display();
  cursorY += 10;
}

// ===== Crash handling =====
void firmwareCrash(const char *reason) {
  firmwareCrashed = true;

  // kill all processes
  for (int i = 0; i < procCount; i++) {
    procs[i].alive = false;
  }

  oledClear();
  display.setCursor(0, 0);
  display.setTextSize(1);
  display.setTextColor(SH110X_WHITE);
  display.print("FIRMWARE CRASHED :(");
  display.setCursor(0, 12);
  display.print(reason);
  display.display();
}

// ===== VM helpers =====
int popVM(Process &p) { return p.stack[p.sp--]; }
void pushVM(Process &p, int v) { p.stack[++p.sp] = v; }

// ===== Syscall handler =====
void handleSyscall(Process &p, int sysNum) {
  switch (sysNum) {
    case SYS_PRINT_INT: {
      int v = popVM(p);
      String line = "P";
      line += p.pid;
      line += ": ";
      line += v;
      oledPrintLine(line);
      break;
    }
    default:
      // unknown syscall, ignore for now
      break;
  }
}

// ===== Single VM step (one or few instructions) =====
void stepProcess(Process &p, int maxInstructions = 8) {
  int executed = 0;

  while (p.alive && executed < maxInstructions && p.ip < p.codeLen) {
    int op = p.code[p.ip++];
    executed++;

    switch (op) {
      case OP_PUSH:
        pushVM(p, p.code[p.ip++]);
        break;

      case OP_ADD: {
        int b = popVM(p);
        int a = popVM(p);
        pushVM(p, a + b);
        break;
      }

      case OP_SUB: {
        int b = popVM(p);
        int a = popVM(p);
        pushVM(p, a - b);
        break;
      }

      case OP_MUL: {
        int b = popVM(p);
        int a = popVM(p);
        pushVM(p, a * b);
        break;
      }

      case OP_DIV: {
        int b = popVM(p);
        int a = popVM(p);
        if (b == 0) {
          firmwareCrash("DIV BY ZERO");
          return;
        }
        pushVM(p, a / b);
        break;
      }

      case OP_PRINT: {
        int v = popVM(p);
        String line = "P";
        line += p.pid;
        line += ": ";
        line += v;
        oledPrintLine(line);
        break;
      }

      case OP_LOAD:
        pushVM(p, p.vars[p.code[p.ip++]]);
        break;

      case OP_STORE:
        p.vars[p.code[p.ip]] = popVM(p);
        p.ip++;
        break;

      case OP_JMP:
        p.ip = p.code[p.ip];
        break;

      case OP_JZ: {
        int addr = p.code[p.ip++];
        if (popVM(p) == 0) p.ip = addr;
        break;
      }

      case OP_SYS: {
        int sysNum = p.code[p.ip++];
        handleSyscall(p, sysNum);
        break;
      }

      case OP_HALT:
        // normal VM halt for that process
        p.alive = false;
        return;

      default:
        // invalid opcode -> kill process
        p.alive = false;
        return;
    }
  }
}

// ===== Scheduler =====
void schedulerTick() {
  if (procCount == 0) return;
  if (firmwareCrashed) return;

  // round-robin
  for (int i = 0; i < procCount; i++) {
    currentProc = (currentProc + 1) % procCount;
    Process &p = procs[currentProc];
    if (!p.alive) continue;
    stepProcess(p);
  }
}

// ===== Example apps (bytecode) =====
// App 1: count up from 1 to 5, printing each via SYS_PRINT_INT
const int app1_code[] = {
  // PUSH 1; STORE 0 (x)
  OP_PUSH, 1,
  OP_STORE, 0,

  // LOAD x; SYS_PRINT_INT
  OP_LOAD, 0,
  OP_SYS, SYS_PRINT_INT,

  // x = x + 1
  OP_LOAD, 0,
  OP_PUSH, 1,
  OP_ADD,
  OP_STORE, 0,

  // LOAD x; SYS_PRINT_INT
  OP_LOAD, 0,
  OP_SYS, SYS_PRINT_INT,

  // x = x + 1
  OP_LOAD, 0,
  OP_PUSH, 1,
  OP_ADD,
  OP_STORE, 0,

  // LOAD x; SYS_PRINT_INT
  OP_LOAD, 0,
  OP_SYS, SYS_PRINT_INT,

  // x = x + 1
  OP_LOAD, 0,
  OP_PUSH, 1,
  OP_ADD,
  OP_STORE, 0,

  // LOAD x; SYS_PRINT_INT
  OP_LOAD, 0,
  OP_SYS, SYS_PRINT_INT,

  // x = x + 1
  OP_LOAD, 0,
  OP_PUSH, 1,
  OP_ADD,
  OP_STORE, 0,

  // LOAD x; SYS_PRINT_INT
  OP_LOAD, 0,
  OP_SYS, SYS_PRINT_INT,

  OP_HALT
};

// App 2: print 100, 90, 80, 70, 60 via SYS_PRINT_INT
const int app2_code[] = {
  OP_PUSH, 100,
  OP_SYS, SYS_PRINT_INT,

  OP_PUSH, 90,
  OP_SYS, SYS_PRINT_INT,

  OP_PUSH, 80,
  OP_SYS, SYS_PRINT_INT,

  OP_PUSH, 70,
  OP_SYS, SYS_PRINT_INT,

  OP_PUSH, 60,
  OP_SYS, SYS_PRINT_INT,

  OP_HALT
};

// ===== Process creation / management =====
int createProcess(const int *code, int len) {
  if (procCount >= MAX_PROCS) return -1;
  Process &p = procs[procCount];
  p.code = code;
  p.codeLen = len;
  p.ip = 0;
  p.sp = -1;
  memset(p.vars, 0, sizeof(p.vars));
  p.alive = true;
  p.pid = nextPid++;
  return procCount++;
}

void killProcessByPid(int pid) {
  for (int i = 0; i < procCount; i++) {
    if (procs[i].pid == pid) {
      procs[i].alive = false;
    }
  }
}

void killAllProcesses() {
  for (int i = 0; i < procCount; i++) {
    procs[i].alive = false;
  }
}

void listProcesses() {
  oledPrintLine("PS:");
  for (int i = 0; i < procCount; i++) {
    String line = "PID ";
    line += procs[i].pid;
    line += " : ";
    line += (procs[i].alive ? "RUN" : "DEAD");
    oledPrintLine(line);
  }
}

// ===== Kernel init / reboot =====
void kernelInit() {
  procCount = 0;
  currentProc = 0;
  nextPid = 1;
  firmwareCrashed = false;
  autoScheduling = true;

  oledClear();
  oledPrintLine("Kernel booted");

  createProcess(app1_code, sizeof(app1_code) / sizeof(int));
  createProcess(app2_code, sizeof(app2_code) / sizeof(int));
}

// ===== Serial command handling =====
String inputLine;

void printHelp() {
  oledPrintLine("HELP:");
  oledPrintLine("INFO, PS, START1, START2");
  oledPrintLine("KILL <pid>, KILLALL");
  oledPrintLine("CLEAR/CLS, ECHO <txt>");
  oledPrintLine("TICK, RUNALL, PAUSE");
  oledPrintLine("HALT, CRASH_MEM");
  oledPrintLine("REBOOT");
}

void handleCommand(const String &cmdRaw) {
  String cmd = cmdRaw;
  cmd.trim();
  if (cmd.length() == 0) return;

  // Uppercase for simple parsing
  String upper = cmd;
  upper.toUpperCase();

  // Simple split: first token + rest
  int spaceIdx = upper.indexOf(' ');
  String op = (spaceIdx == -1) ? upper : upper.substring(0, spaceIdx);
  String arg = (spaceIdx == -1) ? "" : cmd.substring(spaceIdx + 1); // keep original case for ECHO

  if (op == "HELP") {
    printHelp();
  } else if (op == "INFO") {
    oledPrintLine("ESP32 VM OS");
    oledPrintLine("Procs max: 4");
  } else if (op == "PS") {
    listProcesses();
  } else if (op == "START1") {
    int idx = createProcess(app1_code, sizeof(app1_code) / sizeof(int));
    if (idx >= 0) {
      String line = "Started app1 PID ";
      line += procs[idx].pid;
      oledPrintLine(line);
    } else {
      oledPrintLine("No slot for app1");
    }
  } else if (op == "START2") {
    int idx = createProcess(app2_code, sizeof(app2_code) / sizeof(int));
    if (idx >= 0) {
      String line = "Started app2 PID ";
      line += procs[idx].pid;
      oledPrintLine(line);
    } else {
      oledPrintLine("No slot for app2");
    }
  } else if (op == "KILL") {
    if (arg.length() > 0) {
      int pid = arg.toInt();
      killProcessByPid(pid);
      String line = "Killed PID ";
      line += pid;
      oledPrintLine(line);
    } else {
      oledPrintLine("KILL <pid>");
    }
  } else if (op == "KILLALL") {
    killAllProcesses();
    oledPrintLine("All procs killed");
  } else if (op == "CLEAR" || op == "CLS") {
    oledClear();
  } else if (op == "ECHO") {
    if (arg.length() > 0) {
      oledPrintLine(arg);
    } else {
      oledPrintLine("");
    }
  } else if (op == "TICK") {
    schedulerTick();
    oledPrintLine("Tick");
  } else if (op == "RUNALL") {
    autoScheduling = true;
    oledPrintLine("Auto sched ON");
  } else if (op == "PAUSE") {
    autoScheduling = false;
    oledPrintLine("Auto sched OFF");
  } else if (op == "HALT") {
    firmwareCrash("HALT CMD");
  } else if (op == "CRASH_MEM") {
    firmwareCrash("MEM CRASH CMD");
  } else if (op == "REBOOT") {
    kernelInit();
  } else {
    oledPrintLine("Unknown: " + cmd);
  }
}

// ===== Setup / Loop =====
void setup() {
  Serial.begin(115200);
  Wire.begin(I2C_SDA, I2C_SCL);
  display.begin(0x3C, true);
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SH110X_WHITE);
  display.setCursor(0, 0);
  display.display();

  kernelInit();

  Serial.println("ESP32 VM OS ready.");
  Serial.println("Type HELP and press Enter.");
}

void loop() {
  // Serial line reader
  while (Serial.available() > 0) {
    char c = Serial.read();
    if (c == '\r') continue;
    if (c == '\n') {
      handleCommand(inputLine);
      inputLine = "";
    } else {
      inputLine += c;
    }
  }

  // Scheduler
  if (!firmwareCrashed && autoScheduling) {
    schedulerTick();
  }

  delay(50); // crude timeslice
}