AudioEngine.cpp

#include "AudioEngine.h"
#include "NetAudio.h"
#include <M5Unified.h>
#include <esp_heap_caps.h>
#include <cstring>
#include <cstdint>
#include <cstdlib>

AudioEngine Engine;

// ---------------------------------------------------------------------------
// Setup
// ---------------------------------------------------------------------------
bool AudioEngine::begin() {
  const size_t bytes = MAX_SAMPLES * sizeof(int16_t);

  for (uint8_t i = 0; i < NUM_TRACKS; ++i) {
    _tracks[i].buf =
        (int16_t*)heap_caps_malloc(bytes, MALLOC_CAP_SPIRAM);
    if (!_tracks[i].buf) return false;
    memset(_tracks[i].buf, 0, bytes);
  }

  _master = (int16_t*)heap_caps_malloc(bytes, MALLOC_CAP_SPIRAM);
  if (!_master) return false;

  // Start clean: neither mic nor speaker active.
  M5.Speaker.end();
  M5.Mic.end();
  return true;
}

// ---------------------------------------------------------------------------
// Recording  (mic active, speaker off)
// ---------------------------------------------------------------------------
void AudioEngine::resetMasterIfTrack(uint8_t track) {
  // Re-recording the master track redefines the tempo grid.
  if ((int8_t)track == _masterTrack) {
    _masterLoopLen = 0;
    _masterTrack   = -1;
  }
}

void AudioEngine::startRecording(uint8_t track) {
  if (track >= NUM_TRACKS) return;
  if (_state == STATE_PLAYING) stopPlayback();

  resetMasterIfTrack(track);

  M5.Speaker.end();           // free the shared I2S bus

  auto cfg = M5.Mic.config();
  cfg.sample_rate       = SAMPLE_RATE;
  cfg.noise_filter_level = 8; // gentle high-pass, kills DC rumble
  M5.Mic.config(cfg);
  M5.Mic.begin();

  // Discard the mic/PDM + DC-filter settling transient, otherwise every
  // take starts with an audible click/pop at the loop seam.
  static int16_t scratch[REC_CHUNK];
  uint32_t toDiscard = MIC_WARMUP_SAMPLES;
  while (toDiscard) {
    size_t n = (toDiscard < REC_CHUNK) ? toDiscard : REC_CHUNK;
    M5.Mic.record(scratch, n, SAMPLE_RATE, false);
    while (M5.Mic.isRecording()) { M5.delay(1); }
    toDiscard -= n;
  }

  _micSource   = MIC_INTERNAL;
  _recTrack    = track;
  _recPos      = 0;
  _recStartMs  = millis();    // beat 0 of the visual metronome
  _tracks[track].length = 0;  // overwrite from scratch
  _state       = STATE_RECORDING;
}

void AudioEngine::serviceRecording() {
  if (_state != STATE_RECORDING) return;

  Track& t = _tracks[_recTrack];

  // Queue another chunk while there is room in the mic DMA and buffer space.
  if (_recPos < MAX_SAMPLES && M5.Mic.isRecording() != 2) {
    size_t chunk = MAX_SAMPLES - _recPos;
    if (chunk > REC_CHUNK) chunk = REC_CHUNK;
    if (M5.Mic.record(t.buf + _recPos, chunk, SAMPLE_RATE, false)) {
      _recPos += chunk;
    }
  }

  // Track buffer is full -> auto-stop.
  if (_recPos >= MAX_SAMPLES) stopRecording();
}

void AudioEngine::stopRecording() {
  if (_state != STATE_RECORDING) return;

  while (M5.Mic.isRecording()) { M5.delay(1); }  // flush last chunk
  M5.Mic.end();

  finalizeTake();
  _state = STATE_IDLE;
}

// ---------------------------------------------------------------------------
// Remote recording (StickS3 over ESP-NOW) - true overdub: playback keeps
// running, the internal mic is never touched.
// ---------------------------------------------------------------------------
void AudioEngine::startRemoteRecording(uint8_t track) {
  if (track >= NUM_TRACKS) return;
  resetMasterIfTrack(track);

  _micSource  = MIC_REMOTE;
  _recTrack   = track;
  _recPos     = 0;
  _recStartMs = millis();
  _tracks[track].length = 0;
  _state      = STATE_RECORDING;     // playback (if any) is left running
}

void AudioEngine::serviceRemoteRecording(NetAudio& net) {
  if (_state != STATE_RECORDING || _micSource != MIC_REMOTE) return;
  Track& tr = _tracks[_recTrack];
  if (_recPos >= MAX_SAMPLES) { stopRemoteRecording(); return; }
  size_t got = net.pull(tr.buf + _recPos, MAX_SAMPLES - _recPos);
  _recPos += got;
}

void AudioEngine::stopRemoteRecording() {
  if (_state != STATE_RECORDING || _micSource != MIC_REMOTE) return;

  // Drop the tail (KEY1-release transient picked up by the mic).
  if (_recPos > REMOTE_TAIL_TRIM) _recPos -= REMOTE_TAIL_TRIM;
  else                            _recPos  = 0;

  finalizeTake();
  _state = STATE_IDLE;
  startPlayback();                   // play the new mix incl. the fresh take
}

// ---------------------------------------------------------------------------
// Shared take finalisation: onset trim, master quantisation, auto-align, fade.
// ---------------------------------------------------------------------------
void AudioEngine::finalizeTake() {
  Track&   tr      = _tracks[_recTrack];
  int16_t* buf     = tr.buf;
  uint32_t rawLen  = _recPos;
  const bool hadMaster = hasMaster();   // captured before we maybe set one

  // (2) Onset trim: shift content so the first hit sits on sample 0.
  uint32_t onset = findOnset(buf, rawLen);
  if (onset > 0) {
    memmove(buf, buf + onset, (rawLen - onset) * sizeof(int16_t));
  }
  uint32_t lenT = rawLen - onset;

  // (1) Master loop length: first take defines the grid, the rest snap to it.
  uint32_t finalLen;
  if (!hadMaster) {
    finalLen       = lenT;
    _masterLoopLen = lenT;
    _masterTrack   = _recTrack;
  } else {
    uint32_t k = (lenT + _masterLoopLen / 2) / _masterLoopLen;  // round
    if (k < 1) k = 1;
    finalLen = k * _masterLoopLen;
    while (finalLen > MAX_SAMPLES && k > 1) finalLen = (--k) * _masterLoopLen;
    if (finalLen > lenT) {  // pad the tail with silence up to the grid length
      memset(buf + lenT, 0, (finalLen - lenT) * sizeof(int16_t));
    }
  }
  tr.length = finalLen;

  // (5) Auto-align a follower's groove to the master via cross-correlation.
  if (hadMaster && finalLen > 0) {
    int32_t shift = bestAlignShift(buf, finalLen);
    if (shift > 0) rotateLeft(buf, finalLen, (uint32_t)shift);
  }

  // (4) Smooth the loop seam.
  applyFades(buf, finalLen);
}

// ---------------------------------------------------------------------------
// Loop-matching helpers
// ---------------------------------------------------------------------------
uint32_t AudioEngine::findOnset(const int16_t* buf, uint32_t len) const {
  if (len == 0) return 0;

  int32_t peak = 1;
  for (uint32_t i = 0; i < len; ++i) {
    int32_t a = abs((int32_t)buf[i]);
    if (a > peak) peak = a;
  }
  int32_t thr = peak / 8;
  if (thr < ONSET_FLOOR) thr = ONSET_FLOOR;

  uint32_t o = 0;
  while (o < len && abs((int32_t)buf[o]) <= thr) ++o;
  if (o >= len) return 0;                 // silent take -> don't trim

  return (o > ONSET_PREROLL) ? o - ONSET_PREROLL : 0;
}

void AudioEngine::applyFades(int16_t* buf, uint32_t len) const {
  if (len < 2 * FADE_SAMPLES) return;
  for (uint32_t i = 0; i < FADE_SAMPLES; ++i) {
    int32_t g = (int32_t)i;                       // 0..FADE_SAMPLES
    buf[i]           = (int16_t)((int32_t)buf[i]           * g / (int32_t)FADE_SAMPLES);
    buf[len - 1 - i] = (int16_t)((int32_t)buf[len - 1 - i] * g / (int32_t)FADE_SAMPLES);
  }
}

// Find the circular left-shift (0..len) that best lines the new loop up with
// the master mix. Down-sampled and windowed so it stays cheap on the ESP32.
int32_t AudioEngine::bestAlignShift(const int16_t* buf, uint32_t len) const {
  if (_masterTrack < 0 || _masterLoopLen == 0 || len == 0) return 0;

  const int16_t* m    = _tracks[_masterTrack].buf;
  const uint32_t mLen = _masterLoopLen;

  uint32_t cmp = mLen < len ? mLen : len;
  if (cmp > ALIGN_CMP_LEN) cmp = ALIGN_CMP_LEN;

  int64_t best      = INT64_MIN;
  int32_t bestShift = 0;

  for (int32_t d = -(int32_t)ALIGN_WINDOW; d <= (int32_t)ALIGN_WINDOW;
       d += (int32_t)ALIGN_DECIM) {
    int64_t acc = 0;
    for (uint32_t i = 0; i < cmp; i += ALIGN_DECIM) {
      uint32_t fi = ((int32_t)i + d + (int32_t)len) % (int32_t)len;
      acc += (int32_t)buf[fi] * (int32_t)m[i % mLen];
    }
    if (acc > best) { best = acc; bestShift = d; }
  }

  return ((bestShift % (int32_t)len) + (int32_t)len) % (int32_t)len;
}

// In-place circular left rotation via the three-reversal trick (no scratch).
void AudioEngine::rotateLeft(int16_t* buf, uint32_t len, uint32_t shift) const {
  shift %= len;
  if (shift == 0) return;
  auto rev = [&](uint32_t a, uint32_t b) {
    while (a < b) { int16_t t = buf[a]; buf[a] = buf[b]; buf[b] = t; ++a; --b; }
  };
  rev(0, shift - 1);
  rev(shift, len - 1);
  rev(0, len - 1);
}

// ---------------------------------------------------------------------------
// Mixing
// ---------------------------------------------------------------------------
// i-outer / track-inner so per-track effect state (LOFI hold) is trivial and
// no big int32 accumulator buffer is needed. ~4 * masterLen ops -> a few ms.
void AudioEngine::render() {
  _masterLen = 0;
  for (uint8_t t = 0; t < NUM_TRACKS; ++t) {
    if (_tracks[t].hasContent() && !_tracks[t].muted) {
      if (_tracks[t].length > _masterLen) _masterLen = _tracks[t].length;
    }
  }
  if (_masterLen == 0) return;

  int16_t hold[NUM_TRACKS] = {0};  // sample-and-hold state for LOFI

  for (uint32_t i = 0; i < _masterLen; ++i) {
    int32_t sum = 0;
    for (uint8_t t = 0; t < NUM_TRACKS; ++t) {
      if (_tracks[t].hasContent() && !_tracks[t].muted) {
        sum += effectedSample(_tracks[t], i, hold[t]);
      }
    }
    _master[i] = clip16(sum);
  }
}

// One track's effect-processed sample at global index i, looped over its own
// length. `hold` carries the LOFI sample-and-hold state across calls.
int32_t AudioEngine::effectedSample(const Track& tr, uint32_t i,
                                    int16_t& hold) const {
  if (!tr.hasContent()) return 0;

  const uint32_t L = tr.length;
  uint32_t idx = i % L;
  if (tr.fx == FX_REVERSE) idx = (L - 1) - idx;

  int32_t s = tr.buf[idx];
  switch (tr.fx) {
    case FX_BITCRUSH:
      s &= CRUSH_MASK;
      break;
    case FX_LOFI:
      if ((i % LOFI_DECIM) == 0) hold = (int16_t)s;
      s = hold;
      break;
    case FX_DRIVE:
      s = clip16(s * DRIVE_GAIN);
      break;
    case FX_ECHO: {
      uint32_t d = (idx >= ECHO_DELAY) ? idx - ECHO_DELAY : idx + L - ECHO_DELAY;
      s = clip16(s + (tr.buf[d] >> 1));
      break;
    }
    default:  // FX_NONE / FX_REVERSE handled above
      break;
  }
  return s;
}

// ---------------------------------------------------------------------------
// Merge: bake src (with its effect) into dst, then free src. The combined
// audio replaces dst and its effect is reset (already baked in).
// ---------------------------------------------------------------------------
void AudioEngine::mergeTracks(uint8_t dst, uint8_t src) {
  if (dst >= NUM_TRACKS || src >= NUM_TRACKS || dst == src) return;
  Track& a = _tracks[dst];
  Track& b = _tracks[src];
  if (!a.hasContent() && !b.hasContent()) return;

  const uint32_t L = a.hasContent() ? a.length : b.length;
  if (L == 0 || L > MAX_SAMPLES) return;

  const bool wasPlaying = (_state == STATE_PLAYING);
  if (wasPlaying) stopPlayback();      // frees _master to use as scratch

  int16_t holdA = 0, holdB = 0;
  for (uint32_t i = 0; i < L; ++i) {
    int32_t s = effectedSample(a, i, holdA) + effectedSample(b, i, holdB);
    _master[i] = clip16(s);
  }
  memcpy(a.buf, _master, L * sizeof(int16_t));
  a.length = L;
  a.fx     = FX_NONE;                  // effects are now baked into the audio

  // Free the source; keep a tempo grid if the source defined it.
  const bool srcWasMaster = ((int8_t)src == _masterTrack);
  b.length = 0; b.muted = false; b.fx = FX_NONE;
  if (srcWasMaster) { _masterTrack = (int8_t)dst; _masterLoopLen = L; }

  if (wasPlaying) startPlayback();
}

// ---------------------------------------------------------------------------
// Playback  (speaker active, mic off)
// ---------------------------------------------------------------------------
void AudioEngine::startPlayback() {
  if (_state == STATE_RECORDING) return;

  render();
  if (_masterLen == 0) return;  // nothing recorded yet

  M5.Mic.end();                 // free the shared I2S bus
  M5.Speaker.begin();
  M5.Speaker.setVolume(_volume);

  // repeat = 0xFFFFFFFF -> effectively endless loop, handled by the
  // speaker's own background task.
  M5.Speaker.playRaw(_master, _masterLen, SAMPLE_RATE,
                     false /*mono*/, 0xFFFFFFFFu, 0, true);

  _playStartMs = millis();
  _state = STATE_PLAYING;
}

void AudioEngine::stopPlayback() {
  if (_state != STATE_PLAYING) return;
  M5.Speaker.stop();
  M5.Speaker.end();
  _state = STATE_IDLE;
}

void AudioEngine::refreshPlayback() {
  if (_state != STATE_PLAYING) return;
  render();
  if (_masterLen == 0) { stopPlayback(); return; }
  M5.Speaker.playRaw(_master, _masterLen, SAMPLE_RATE,
                     false, 0xFFFFFFFFu, 0, true);
  _playStartMs = millis();
}

// ---------------------------------------------------------------------------
// Editing
// ---------------------------------------------------------------------------
void AudioEngine::cycleFx(uint8_t track) {
  if (track >= NUM_TRACKS) return;
  _tracks[track].fx = (Fx)((_tracks[track].fx + 1) % FX_COUNT);
  refreshPlayback();
}

void AudioEngine::toggleMute(uint8_t track) {
  if (track >= NUM_TRACKS) return;
  _tracks[track].muted = !_tracks[track].muted;
  refreshPlayback();
}

void AudioEngine::setVolume(uint8_t v) {
  _volume = v;
  M5.Speaker.setVolume(_volume);  // master volume, safe to set any time
}

void AudioEngine::clearTrack(uint8_t track) {
  if (track >= NUM_TRACKS) return;
  _tracks[track].length = 0;
  _tracks[track].muted  = false;
  _tracks[track].fx     = FX_NONE;

  // Clearing the master track drops the tempo grid; the next take redefines it.
  if (track == _masterTrack) {
    _masterLoopLen = 0;
    _masterTrack   = -1;
  }
  refreshPlayback();
}

// ---------------------------------------------------------------------------
// Progress reporting for the UI
// ---------------------------------------------------------------------------
float AudioEngine::progress() const {
  if (_state == STATE_RECORDING) {
    return (float)_recPos / (float)MAX_SAMPLES;
  }
  if (_state == STATE_PLAYING && _masterLen) {
    float loopMs = 1000.0f * _masterLen / SAMPLE_RATE;
    float pos    = fmodf((float)(millis() - _playStartMs), loopMs);
    return pos / loopMs;
  }
  return 0.0f;
}