#include <SPI.h>

/*
  main.ino
  Злиття вашого початкового коду (пункт 1) + UART/dispatcher/команди (пункти 2,3).
*/

// ---------------- HW mapping (з вашого початкового коду) ----------------
constexpr uint8_t LATCH_PIN = PB0; // STCP (latch)
constexpr uint8_t BTN_PIN   = PB3; // кнопка B1 -> GND

// Polarity flags (як у вас)
constexpr bool INV_SEG  = true;   // segments are active-LOW
constexpr bool INV_DIG  = false;  // digits active-HIGH
constexpr bool SWAP_SEND = false; // keep false for chain MCU->srSeg->srDig

// 7-seg map for CC (A..G,DP = 1 is ON); we invert at runtime if INV_SEG=true
const uint8_t DIGIT_MAP[10] = {
  0b00111111, // 0
  0b00000110, // 1
  0b01011011, // 2
  0b01001111, // 3
  0b01100110, // 4
  0b01101101, // 5
  0b01111101, // 6
  0b00000111, // 7
  0b01111111, // 8
  0b01101111  // 9
};

inline uint8_t segFix(uint8_t v) { return INV_SEG ? ~v : v; }
inline uint8_t digFix(uint8_t v) { return INV_DIG ? ~v : v; }

inline void shift16(uint8_t first, uint8_t second) {
  digitalWrite(LATCH_PIN, LOW);
  if (SWAP_SEND) { SPI.transfer(second); SPI.transfer(first); }
  else           { SPI.transfer(first);  SPI.transfer(second); }
  digitalWrite(LATCH_PIN, HIGH);
}

inline void allDigitsOff() { shift16(digFix(0x00), 0x00); } // blank

// ---------------- shared state ----------------
volatile bool countDown = false; // default: up
volatile int counter = 0;        // 0..9999
volatile int stepSize = 1;       // 1..99

// timing for debounce and stepping
uint32_t lastBtnMs = 0, lastStepMs = 0;

// ---------------- button reading (debounced toggle) ----------------
bool readButtonToggle() {
  const uint32_t now = millis();
  static bool prev = HIGH;
  bool cur = digitalRead(BTN_PIN);
  if (cur != prev && (now - lastBtnMs > 50)) {
    lastBtnMs = now; prev = cur;
    return (cur == LOW); // pressed
  }
  return false;
}

// ---------------- display: multiplex whole number ----------------
void displayNumber(int value) {
  if (value < 0) value = (10000 + (value % 10000)) % 10000;
  value %= 10000;

  uint8_t d[4] = {
    (uint8_t)(value / 1000),
    (uint8_t)((value / 100) % 10),
    (uint8_t)((value / 10) % 10),
    (uint8_t)(value % 10)
  };

  for (uint8_t i = 0; i < 4; i++) {
    allDigitsOff(); // anti-ghosting
    uint8_t seg = segFix(DIGIT_MAP[d[i]]);
    uint8_t dig = digFix(1 << i); // Q0..Q3 -> DIG1..DIG4
    shift16(dig, seg);
    delayMicroseconds(800); // small on-time per digit
  }
}

// ---------------- Dispatcher (cooperative) ----------------
// Dispatcher (cooperative)
typedef void (*TaskFn)(void);

struct Task {
  TaskFn fn;
  uint32_t periodMs;
  uint32_t lastRun;
};

#define NUM_TASKS 4
Task tasks[NUM_TASKS];

void scheduleTask(uint8_t idx, TaskFn f, uint32_t period) {
  if (idx >= NUM_TASKS) return;
  tasks[idx].fn = f;
  tasks[idx].periodMs = period;
  tasks[idx].lastRun = millis();
}

// ---------------- UART parser & helpers ----------------
String serialBuf = "";

// trim
String trimStr(const String &s) {
  int i = 0, j = s.length()-1;
  while (i <= j && isspace(s[i])) i++;
  while (j >= i && isspace(s[j])) j--;
  if (i==0 && j==s.length()-1) return s;
  return s.substring(i, j+1);
}

String toUpperStr(const String &s) {
  String r = s;
  r.toUpperCase();
  return r;
}

// simple tokenize (space-separated)
void tokenize(const String &s, String tokens[], int &count, int maxTokens=6) {
  count = 0;
  String cur = "";
  bool inQuote = false;
  for (int i=0;i<s.length();++i) {
    char c = s[i];
    if (c == '"') { inQuote = !inQuote; continue; }
    if (!inQuote && isspace(c)) {
      if (cur.length()>0) {
        if (count < maxTokens) tokens[count++] = cur;
        cur = "";
      }
    } else cur += c;
  }
  if (cur.length()>0 && count < maxTokens) tokens[count++] = cur;
}

// send
void sendOK(const String &msg) { Serial.print("OK "); Serial.println(msg); }
void sendERR(const String &msg) { Serial.print("ERR "); Serial.println(msg); }

// parse integer
bool parseIntStr(const String &s, long &val) {
  char *endptr;
  val = strtol(s.c_str(), &endptr, 10);
  if (*endptr != 0) return false;
  return true;
}

// process one command line (case-insensitive)
void processCommandLine(String line) {
  line = trimStr(line);
  if (line.length() == 0) return;

  String tokens[6];
  int tcount = 0;
  tokenize(line, tokens, tcount, 6);
  if (tcount == 0) return;

  String cmd = toUpperStr(tokens[0]); // GET, RESET, SET

  if (cmd != "GET" && cmd != "RESET" && cmd != "SET") {
    sendERR("Unknown command");
    return;
  }

  // default param: -A
  String param = "-A";
  String arg = "";

  if (tcount >= 2) param = toUpperStr(tokens[1]);
  if (tcount >= 3) arg = tokens[2];

  // normalize param
  if (!(param.length() >= 2 && param[0] == '-')) param = "-" + toUpperStr(param);

  // GET
  if (cmd == "GET") {
    if (param == "-C") {
      Serial.print("COUNTER ");
      Serial.println(counter);
    } else if (param == "-D") {
      Serial.print("DIRECTION ");
      Serial.println(countDown ? "DOWN" : "UP");
    } else if (param == "-S") {
      Serial.print("STEP ");
      Serial.println(stepSize);
    } else { // -A or other
      Serial.print("COUNTER ");
      Serial.println(counter);
      Serial.print("DIRECTION ");
      Serial.println(countDown ? "DOWN" : "UP");
      Serial.print("STEP ");
      Serial.println(stepSize);
    }
    return;
  }

  // RESET
  if (cmd == "RESET") {
    if (param == "-C") {
      counter = 0;
      sendOK("Counter reset to 0");
    } else if (param == "-D") {
      countDown = false;
      sendOK("Direction set to UP");
    } else if (param == "-S") {
      stepSize = 1;
      sendOK("Step reset to 1");
    } else { // -A
      counter = 0;
      countDown = false;
      stepSize = 1;
      sendOK("All reset (counter=0, direction=UP, step=1)");
    }
    return;
  }

  // SET
  if (cmd == "SET") {
    if (param == "-C") {
      if (arg.length() == 0) { sendERR("SET -C requires value"); return; }
      long v;
      if (!parseIntStr(arg, v)) { sendERR("Invalid number"); return; }
      if (v < 0 || v > 9999) { sendERR("Counter out of range (0..9999)"); return; }
      counter = (int)v;
      sendOK("Counter set");
      return;
    } else if (param == "-S") {
      if (arg.length() == 0) { sendERR("SET -S requires value"); return; }
      long v;
      if (!parseIntStr(arg, v)) { sendERR("Invalid number"); return; }
      if (v < 1 || v > 99) { sendERR("Step out of range (1..99)"); return; }
      stepSize = (int)v;
      sendOK("Step set");
      return;
    } else if (param == "-D") {
      if (arg.length() == 0) { sendERR("SET -D requires UP or DOWN"); return; }
      String up = toUpperStr(arg);
      if (up == "UP" || up == "+") {
        countDown = false;
        sendOK("Direction set to UP");
        return;
      } else if (up == "DOWN" || up == "-") {
        countDown = true;
        sendOK("Direction set to DOWN");
        return;
      } else {
        sendERR("Direction must be UP or DOWN (or + / -)");
        return;
      }
    } else {
      sendERR("SET unknown parameter");
      return;
    }
  }
}

// UART task: build lines, call processor
void taskUART() {
  while (Serial.available()) {
    char c = Serial.read();
    if (c == '\r') continue;
    if (c == '\n') {
      if (serialBuf.length() > 0) {
        processCommandLine(serialBuf);
        serialBuf = "";
      }
    } else {
      serialBuf += c;
      if (serialBuf.length() > 200) {
        serialBuf = "";
        sendERR("Input too long");
      }
    }
  }
}

// counter stepping task (uses stepSize)
void taskCounter() {
  uint32_t now = millis();
  if (now - lastStepMs >= 500) { // 500 ms per original
    lastStepMs = now;
    if (countDown) counter -= stepSize;
    else counter += stepSize;

    if (counter > 9999) counter = counter % 10000;
    if (counter < 0) counter = ((counter % 10000) + 10000) % 10000;
  }
}

// button task: toggle on press (using readButtonToggle)
void taskButton() {
  if (readButtonToggle()) {
    countDown = !countDown;
    Serial.print("BUTTON: direction=");
    Serial.println(countDown ? "DOWN" : "UP");
  }
}

// display task: call multiplex routine (this is a small blocking call)
void taskDisplay() {
  displayNumber(counter);
}

// ---------------- setup & loop ----------------
void setup() {
  pinMode(LATCH_PIN, OUTPUT);
  pinMode(BTN_PIN, INPUT_PULLUP);
  SPI.begin();
  // SPI pins: SCK=PA5, MOSI=PA7 set by SPI.begin(); keep settings similar to original
  SPI.beginTransaction(SPISettings(8000000, MSBFIRST, SPI_MODE0));
  allDigitsOff();

  Serial.begin(115200);
  delay(20);
  Serial.println("READY");

  // defaults
  countDown = false;
  counter = 0;
  stepSize = 1;
  lastStepMs = millis();

  // schedule tasks
  scheduleTask(0, taskUART, 10);    // UART polling
  scheduleTask(1, taskButton, 20);  // button polling
  scheduleTask(2, taskCounter, 50); // counter (internal 500ms tick)
  scheduleTask(3, taskDisplay, 4);  // display multiplex cycles
}

void loop() {
  uint32_t now = millis();
  for (int i = 0; i < NUM_TASKS; ++i) {
    if ((now - tasks[i].lastRun) >= tasks[i].periodMs) {
      tasks[i].lastRun = now;
      if (tasks[i].fn) tasks[i].fn();
    }
  }
}