#include <Arduino.h>
#include <string.h>
#include <EEPROM.h>
enum Cmd {
  CMD_HELP, CMD_INSERT, CMD_DELETE, CMD_SEARCH,
  CMD_PRINT_IN, CMD_PRINT_PRE, CMD_PRINT_POST,
  CMD_ADDR, CMD_MEM, CMD_CLEAR,
  CMD_SAVE, CMD_LOAD, CMD_EERASE,
  CMD_UNKNOWN
};
// =========================================================
// Forward declarations (so phases can be ordered for teaching)
// =========================================================
struct Node;
extern Node* root;

bool insert(Node*& r, int key);        // Phase #3
Node* findNode(Node* r, int key);      // Phase #3
bool removeKey(Node*& r, int key);     // Phase #7

void clearPool();                      // Phase #2
uint8_t usedCount();                   // Phase #2

void printInorder(Node* r);            // Phase #4
void printPreorder(Node* r);           // Phase #4
void printPostorder(Node* r);          // Phase #4
void printWithAddresses(Node* r);      // Phase #4

// =========================================================
// #1 PHASE 1: "FROM THE START" — SERIAL COMMANDS + SAVE/LOAD
//     Goal: immediately show the tree + store/read it.
// =========================================================

// EEPROM layout (simple + robust):
// addr 0: magic byte (0x42 means "valid data")
// addr 1: count (0..MAX_NODES)
// addr 2.. : int16_t keys in preorder (2 bytes each)
const uint8_t EEPROM_MAGIC = 0x42;
const int EEPROM_ADDR_MAGIC = 0;
const int EEPROM_ADDR_COUNT = 1;
const int EEPROM_ADDR_KEYS  = 2;

// Reads one line from Serial into buf (until '\n').
bool readLine(char* buf, size_t len) {
  if (!Serial.available()) return false;
  size_t n = Serial.readBytesUntil('\n', buf, len - 1);
  buf[n] = '\0';

  // Remove CR (Windows Serial Monitor)
  size_t L = strlen(buf);
  if (L > 0 && buf[L - 1] == '\r') buf[L - 1] = '\0';
  return true;
}

// Basic help text for students (F() keeps text in Flash on AVR boards).
void printHelp() {
  Serial.println();
  Serial.println(F("Commands:"));
  Serial.println(F("  h        : help"));
  Serial.println(F("  i <n>    : insert number n"));
  Serial.println(F("  d <n>    : delete number n   (Phase #7)"));
  Serial.println(F("  s <n>    : search number n"));
  Serial.println(F("  p        : print inorder"));
  Serial.println(F("  pre      : print preorder"));
  Serial.println(F("  post     : print postorder"));
  Serial.println(F("  a        : show addresses (pointers)"));
  Serial.println(F("  m        : memory usage (pool)"));
  Serial.println(F("  c        : clear tree (RAM only)"));
  Serial.println(F("  w        : write/save tree to EEPROM"));
  Serial.println(F("  r        : read/load tree from EEPROM"));
  Serial.println(F("  x        : erase EEPROM"));
  Serial.println();
}

// Collect preorder keys into an array (for saving).
// (We store keys only; pointers cannot be stored and reused.)
void collectPreorder(Node* r, int16_t* out, uint8_t& idx, uint8_t maxCount);

// Save current tree into EEPROM.
void eepromSaveTree(Node* r, uint8_t maxCount) {
  int16_t keys[255];          // upper bound, but we will only fill up to maxCount
  uint8_t count = 0;

  // Safer stack use on small boards: only allocate what we need.
  // For UNO teaching, keep MAX_NODES small; this array is "over", but idx is limited.
  // If you want the stack minimal, set keys[MAX_NODES] below and keep MAX_NODES small.
  collectPreorder(r, keys, count, maxCount);

  EEPROM.update(EEPROM_ADDR_MAGIC, EEPROM_MAGIC);
  EEPROM.update(EEPROM_ADDR_COUNT, count);

  int addr = EEPROM_ADDR_KEYS;
  for (uint8_t i = 0; i < count; i++) {
    EEPROM.put(addr, keys[i]);                 // 2 bytes per key
    addr += (int)sizeof(int16_t);
  }

  Serial.print(F("Saved "));
  Serial.print(count);
  Serial.println(F(" keys to EEPROM (preorder)."));
}

// Load tree from EEPROM (rebuild by inserting keys in stored order).
void eepromLoadTree(Node*& r, uint8_t maxCount) {
  uint8_t magic = EEPROM.read(EEPROM_ADDR_MAGIC);
  if (magic != EEPROM_MAGIC) {
    Serial.println(F("No valid EEPROM data found."));
    return;
  }

  uint8_t count = EEPROM.read(EEPROM_ADDR_COUNT);
  if (count > maxCount) count = maxCount;

  clearPool();
  r = nullptr;

  int addr = EEPROM_ADDR_KEYS;
  for (uint8_t i = 0; i < count; i++) {
    int16_t key;
    EEPROM.get(addr, key);
    addr += (int)sizeof(int16_t);
    insert(r, (int)key);
  }

  Serial.print(F("Loaded "));
  Serial.print(count);
  Serial.println(F(" keys from EEPROM and rebuilt tree."));
}

// Erase EEPROM "valid flag" (tree becomes "not saved").
void eepromErase() {
  EEPROM.update(EEPROM_ADDR_MAGIC, 0x00);
  EEPROM.update(EEPROM_ADDR_COUNT, 0x00);
  Serial.println(F("EEPROM erased (magic cleared)."));
}

// =========================================================
// #2 PHASE 2: DATA STRUCTURE (POINTERS) + FIXED MEMORY POOL
// =========================================================

struct Node {
  int key;
  Node* left;
  Node* right;
};

// Keep this small for UNO; increase on boards with more SRAM.
const uint8_t MAX_NODES = 25;

Node pool[MAX_NODES];
bool used[MAX_NODES];
Node* root = nullptr;

// Allocates a node from the fixed pool (no heap fragmentation).
Node* allocNode(int key) {
  for (uint8_t i = 0; i < MAX_NODES; i++) {
    if (!used[i]) {
      used[i] = true;
      pool[i].key = key;
      pool[i].left = nullptr;
      pool[i].right = nullptr;
      return &pool[i];
    }
  }
  return nullptr; // pool full
}

// Frees a node back into the pool (so it can be reused).
bool freeNode(Node* n) {
  if (!n) return false;
  if (n < &pool[0] || n >= &pool[MAX_NODES]) return false;

  ptrdiff_t idx = n - &pool[0];
  if (idx < 0 || idx >= (ptrdiff_t)MAX_NODES) return false;

  used[(uint8_t)idx] = false;

  // Optional cleanup (nice for debugging)
  pool[(uint8_t)idx].left = nullptr;
  pool[(uint8_t)idx].right = nullptr;
  pool[(uint8_t)idx].key = 0;

  return true;
}

// Resets pool usage flags (tree nodes become "free" again).
void clearPool() {
  for (uint8_t i = 0; i < MAX_NODES; i++) used[i] = false;
}

// Count how many nodes are currently used.
uint8_t usedCount() {
  uint8_t c = 0;
  for (uint8_t i = 0; i < MAX_NODES; i++) if (used[i]) c++;
  return c;
}

// =========================================================
// #3 PHASE 3: BST CORE (REFERENCES & POINTERS)
// =========================================================

// Insert uses Node*& so the function can MODIFY the pointer itself.
// This is the key "reference-to-pointer" lesson.
bool insert(Node*& r, int key) {
  if (r == nullptr) {
    r = allocNode(key);
    return (r != nullptr);
  }
  if (key == r->key) return false; // no duplicates in this demo
  if (key < r->key) return insert(r->left, key);
  return insert(r->right, key);
}

Node* findNode(Node* r, int key) {
  while (r != nullptr) {
    if (key == r->key) return r;
    r = (key < r->key) ? r->left : r->right;
  }
  return nullptr;
}

// =========================================================
// #7 PHASE 7: DELETE / REMOVE (POINTER PRACTICE!)
//     Removes a key from the BST and frees the node to the pool.
// =========================================================

Node* minNode(Node* r) {
  if (!r) return nullptr;
  while (r->left) r = r->left;
  return r;
}

bool removeKey(Node*& r, int key) {
  if (!r) return false;

  if (key < r->key) return removeKey(r->left, key);
  if (key > r->key) return removeKey(r->right, key);

  // Found the node to delete
  Node* toDelete = r;

  // Case 1: no children
  if (!r->left && !r->right) {
    r = nullptr;
    freeNode(toDelete);
    return true;
  }

  // Case 2: only right child
  if (!r->left) {
    r = r->right;
    freeNode(toDelete);
    return true;
  }

  // Case 3: only left child
  if (!r->right) {
    r = r->left;
    freeNode(toDelete);
    return true;
  }

  // Case 4: two children
  Node* succ = minNode(r->right);
  int succKey = succ->key;

  r->key = succKey;
  return removeKey(r->right, succKey);
}

// =========================================================
// #4 PHASE 4: SHOW THE TREE (TRAVERSALS + DEBUG ADDRESSES)
// =========================================================

void printInorder(Node* r) {
  if (!r) return;
  printInorder(r->left);
  Serial.print(r->key);
  Serial.print(' ');
  printInorder(r->right);
}

void printPreorder(Node* r) {
  if (!r) return;
  Serial.print(r->key);
  Serial.print(' ');
  printPreorder(r->left);
  printPreorder(r->right);
}

void printPostorder(Node* r) {
  if (!r) return;
  printPostorder(r->left);
  printPostorder(r->right);
  Serial.print(r->key);
  Serial.print(' ');
}

// Makes pointers visible: prints node address + left/right addresses.
void printWithAddresses(Node* r) {
  if (!r) return;

  Serial.print(F("key="));
  Serial.print(r->key);

  Serial.print(F(" @0x"));
  Serial.print((uintptr_t)r, HEX);

  Serial.print(F(" L=0x"));
  Serial.print((uintptr_t)r->left, HEX);

  Serial.print(F(" R=0x"));
  Serial.println((uintptr_t)r->right, HEX);

  printWithAddresses(r->left);
  printWithAddresses(r->right);
}

// Preorder collector for EEPROM save (keys only)
void collectPreorder(Node* r, int16_t* out, uint8_t& idx, uint8_t maxCount) {
  if (!r || idx >= maxCount) return;
  out[idx++] = (int16_t)r->key;
  collectPreorder(r->left, out, idx, maxCount);
  collectPreorder(r->right, out, idx, maxCount);
}

// =========================================================
// #5 PHASE 5: PARSING COMMANDS + SWITCH/CASE DISPATCH
// =========================================================



Cmd parseCmd(const char* c) {
  if (!c || !*c) return CMD_UNKNOWN;

  // Single-letter commands
  if (c[1] == '\0') {
    switch (c[0]) {
      case 'h': return CMD_HELP;
      case 'i': return CMD_INSERT;
      case 'd': return CMD_DELETE;
      case 's': return CMD_SEARCH;
      case 'p': return CMD_PRINT_IN;
      case 'a': return CMD_ADDR;
      case 'm': return CMD_MEM;
      case 'c': return CMD_CLEAR;
      case 'w': return CMD_SAVE;
      case 'r': return CMD_LOAD;
      case 'x': return CMD_EERASE;
      default:  return CMD_UNKNOWN;
    }
  }

  // Multi-letter commands (kept simple for teaching)
  if (strcmp(c, "pre") == 0)  return CMD_PRINT_PRE;
  if (strcmp(c, "post") == 0) return CMD_PRINT_POST;

  return CMD_UNKNOWN;
}

// =========================================================
// #6 PHASE 6: PROGRAM FLOW (SETUP + LOOP)
// =========================================================

void setup() {
  Serial.begin(115200);
  clearPool();
  root = nullptr;

  Serial.println(F("Arduino BST Pointer/Reference Demo (EEPROM save/load + delete)"));
  Serial.println(F("Type 'h' for help."));
}

void loop() {
  static char line[48];
  if (!readLine(line, sizeof(line))) return;

  char* cmd = strtok(line, " ");
  char* arg = strtok(nullptr, " ");
  Cmd c = parseCmd(cmd);

  switch (c) {
    case CMD_HELP:
      printHelp();
      break;

    case CMD_INSERT: {
      if (!arg) { Serial.println(F("Missing number. Example: i 42")); break; }
      int x = atoi(arg);
      bool ok = insert(root, x);
      Serial.println(ok ? F("OK.") : F("Insert failed (duplicate or pool full)."));
      break;
    }

    case CMD_DELETE: {
      if (!arg) { Serial.println(F("Missing number. Example: d 42")); break; }
      int x = atoi(arg);
      bool ok = removeKey(root, x);
      Serial.println(ok ? F("Deleted.") : F("Not found (nothing deleted)."));
      break;
    }

    case CMD_SEARCH: {
      if (!arg) { Serial.println(F("Missing number. Example: s 42")); break; }
      int x = atoi(arg);
      Node* n = findNode(root, x);
      if (!n) {
        Serial.println(F("Not found."));
      } else {
        Serial.print(F("Found: key="));
        Serial.print(n->key);
        Serial.print(F(" @0x"));
        Serial.println((uintptr_t)n, HEX);
      }
      break;
    }

    case CMD_PRINT_IN:
      printInorder(root); Serial.println();
      break;

    case CMD_PRINT_PRE:
      printPreorder(root); Serial.println();
      break;

    case CMD_PRINT_POST:
      printPostorder(root); Serial.println();
      break;

    case CMD_ADDR:
      if (!root) Serial.println(F("Tree empty."));
      else printWithAddresses(root);
      break;

    case CMD_MEM:
      Serial.print(F("Nodes used: "));
      Serial.print(usedCount());
      Serial.print(F("/"));
      Serial.println(MAX_NODES);
      break;

    case CMD_CLEAR:
      clearPool();
      root = nullptr;
      Serial.println(F("Tree cleared (RAM)."));
      break;

    case CMD_SAVE:
      eepromSaveTree(root, MAX_NODES);
      break;

    case CMD_LOAD:
      eepromLoadTree(root, MAX_NODES);
      break;

    case CMD_EERASE:
      eepromErase();
      break;

    default:
      Serial.println(F("Unknown command. Type 'h'."));
      break;
  }
}