/*
 * Copyright 2023 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "soundstream.h"

#include <fail-simulator-on-error.h>
#include <thread.h>

#include <debug.hh>

#include "hw/top_matcha/sw/autogen/top_matcha.h"

/// Expose debugging features unconditionally for this compartment.
using Debug = ConditionalDebug<true, "SOUNDSTREAM">;

#define ABS(x) (x > 0 ? x : -x)

// #define kSamples (5 * 16000)
#define kSamples (2000)  // NB: reduce sample count for slow renode
#define kFilterSamples (256)

// Marker symbols to hook to enable/disable profiling
extern "C" {
void __attribute__((noinline)) stats_enable(void) { asm(""); }
void __attribute__((noinline)) stats_disable(void) { asm(""); }
}

void __cheri_compartment("soundstream") entry(void) {
  Debug::log("soundstream (Thread {})", thread_id_get());

  i2s_init();
  ml_top_init();
#if DEVICE_EXISTS_mailbox
  mailbox_init();
#endif

  i2s_rxfifo_clear();
  i2s_irq_acknowledge_all();
  i2s_irq_set_enabled(kI2sIrqRxWatermark, /*enabled=*/true);

  Debug::log("Loading Kelvin binary from SPI");
  ml_top_load_file_from_tar("kelvin.bin");

  // NB: the stack is 4KiB so this uses 1/4 of it
  int16_t samples_left[kFilterSamples] = {0};
  int16_t samples_right[kFilterSamples] = {0};
  size_t index_left = 0;
  size_t index_right = 0;
  int32_t total_left = 0;
  int32_t total_right = 0;

  static int32_t samples[kSamples];
  memset(samples, 0xa5, sizeof(int32_t) * kSamples);

  Debug::log("Setup complete");

  stats_enable();
// NB: once around the loop for collecting profile data
#if 0
  while (true) {
#else
  {
#endif
#if DEVICE_EXISTS_mailbox
    // TODO(sleffler): need custom security core code running and
    //   a way to toggle the gpio associated with the button; for now
    //   just force it to appear as though the button has been pressed.
    mailbox_set_button_pressed(true);

    // Wait until the record switch is pushed
    Debug::log("Wait for button press...");
    mailbox_wait_for_button_pressed();

    mailbox_set_led(/*enabled=*/true);
#else
#endif
    Debug::log("Start recording (max {} samples)...", kSamples);
    i2s_record_begin();

    // Record until our buffer is full or the switch is released.
    int sample = 0;
    while (sample < kSamples
#if DEVICE_EXISTS_mailbox
           && mailbox_button_pressed()
#else
#endif
    ) {
      i2s_irq_set_enabled(kI2sIrqRxWatermark, /*enabled=*/true);

      i2s_wait_for_rx_watermark();

      while (!i2s_rxfifo_is_empty()) {
        uint32_t reg_val = i2s_get_rdata();
        // For each sample, split into left and right channel values,
        // and update the moving average.
        int16_t left = reg_val >> 16;
        int16_t right = reg_val & 0xFFFF;
        total_left -= samples_left[index_left];
        total_right -= samples_right[index_right];
        total_left += left;
        samples_left[index_left] = left;
        total_right += right;
        samples_right[index_right] = right;
        index_left = (index_left + 1) % kFilterSamples;
        index_right = (index_right + 1) % kFilterSamples;
        int16_t mean_left = total_left / kFilterSamples;
        int16_t mean_right = total_right / kFilterSamples;

        // Subtract the moving average from each channel, and repack into the
        // sample buffer.
        uint32_t offset_reg_val = (((left - mean_left) & 0xFFFF) << 16) |
                                  ((right - mean_right) & 0xFFFF);
        samples[sample++] = offset_reg_val;
        if (sample == kSamples
#if DEVICE_EXISTS_mailbox
            || !mailbox_button_pressed()
#else
#endif
        ) {
          break;
        }
      }
    }

    i2s_record_end();
#if DEVICE_EXISTS_mailbox
    mailbox_set_led(/*enabled=*/false);
#else
#endif

    Debug::log("Done recording {} samples", sample);
    int samples_captured = sample;

    // Calculate the min/max of the audio, after correcting DC offsets
    int32_t max = INT16_MIN;
    int32_t min = INT16_MAX;
    for (int i = 1; i < (samples_captured * 2); i += 2) {
      int16_t* samples_s16 = reinterpret_cast<int16_t*>(samples);
      int32_t sample = samples_s16[i];
      if (sample < min) {
        min = sample;
      }
      if (sample > max) {
        max = sample;
      }
    }

    // Calculate a scaling factor, and use this to scale the waveform
    // to a peak of 75% amplitude.
    int32_t scale_max = ((int32_t)max * 100) / ((int32_t)INT16_MAX);
    int32_t scale_min = ABS(((int32_t)min * 100) / ((int32_t)INT16_MIN));
    int32_t scale = scale_max > scale_min ? scale_max : scale_min;
    for (int i = 1; i < (samples_captured * 2); i += 2) {
      int16_t* samples_s16 = reinterpret_cast<int16_t*>(samples);
      int16_t sample = samples_s16[i];
      int16_t scaled_sample = (int16_t)(((int32_t)sample * 100) / scale);
      scaled_sample = (((int32_t)scaled_sample * 75) / 100);
      samples_s16[i] = scaled_sample;
    }

    Debug::log("Processing recorded audio...");

    // 320 x int16
    int iterations_to_process = samples_captured / 320;
    int16_t process_buffer[320];
    int16_t result_buffer[64];
    char result_buffer_encoded[ENCODE_OUT_SIZE(sizeof(result_buffer))];

    struct output_header header;
    memset(&header, 0, sizeof(header));  // NB: resume_pc = 0

    for (int i = 0; i < iterations_to_process; ++i) {
      Debug::log("Iteration {}", i);
      int16_t* samples_s16 = reinterpret_cast<int16_t*>(samples);
      // Extract left channel audio
      for (int j = 0; j < 320; ++j) {
        process_buffer[j] = samples_s16[(i * 320 * 2) + (j * 2) + 1];
      }
      ml_top_set_input(process_buffer, sizeof(process_buffer));

      // Start/resume kelvin
      (void) ml_top_finish_done();  // NB: reset state
      ml_top_resume_ctrl_en(header.resume_pc);

      ml_top_wait_for_finish();

      ml_top_get_output_header(&header);
      Debug::Assert(header.length == sizeof(result_buffer),
                    "Unexpected ML result size");
      ml_top_get_output_data(&header, result_buffer);

      encode((const unsigned char*)result_buffer, sizeof(result_buffer),
             result_buffer_encoded);
      Debug::log("[sound]::ENCODER:{}",
        std::string_view(result_buffer_encoded, sizeof(result_buffer_encoded)));
    }

    Debug::log("[sound]::ENCODER: done");
    Debug::log("Done with processing.");
  }
  stats_disable();
  simulation_exit(0);
}