#include "audio_resampler.h"

#ifdef USE_ESP32

#include "esphome/core/hal.h"

#include <cstring>

namespace esphome::audio {

static const uint32_t READ_WRITE_TIMEOUT_MS = 20;

AudioResampler::AudioResampler(size_t input_buffer_size, size_t output_buffer_size)
    : input_buffer_size_(input_buffer_size), output_buffer_size_(output_buffer_size) {
  this->output_transfer_buffer_ = AudioSinkTransferBuffer::create(output_buffer_size);
}

esp_err_t AudioResampler::add_source(std::weak_ptr<ring_buffer::RingBuffer> &input_ring_buffer) {
  // The zero-copy RingBufferAudioSource is created lazily on the first resample() call, once both the ring
  // buffer (stored here) and the input stream info (set by start()) are available, in either order.
  this->source_ring_buffer_ = input_ring_buffer.lock();
  if (this->source_ring_buffer_ == nullptr) {
    return ESP_ERR_INVALID_STATE;
  }
  return ESP_OK;
}

esp_err_t AudioResampler::add_sink(std::weak_ptr<ring_buffer::RingBuffer> &output_ring_buffer) {
  if (this->output_transfer_buffer_ != nullptr) {
    this->output_transfer_buffer_->set_sink(output_ring_buffer);
    return ESP_OK;
  }
  return ESP_ERR_NO_MEM;
}

#ifdef USE_SPEAKER
esp_err_t AudioResampler::add_sink(speaker::Speaker *speaker) {
  if (this->output_transfer_buffer_ != nullptr) {
    this->output_transfer_buffer_->set_sink(speaker);
    return ESP_OK;
  }
  return ESP_ERR_NO_MEM;
}
#endif

esp_err_t AudioResampler::start(AudioStreamInfo &input_stream_info, AudioStreamInfo &output_stream_info,
                                uint16_t number_of_taps, uint16_t number_of_filters) {
  this->input_stream_info_ = input_stream_info;
  this->output_stream_info_ = output_stream_info;

  if (this->output_transfer_buffer_ == nullptr) {
    return ESP_ERR_NO_MEM;
  }

  if ((input_stream_info.get_bits_per_sample() > 32) || (output_stream_info.get_bits_per_sample() > 32) ||
      (input_stream_info_.get_channels() != output_stream_info.get_channels())) {
    return ESP_ERR_NOT_SUPPORTED;
  }

  // Reject frame sizes that can't be used as the zero-copy source's alignment up front, where the caller checks
  // the return code. The lazy create() in resample() keeps its own guard since it runs before the uint8_t cast.
  const size_t bytes_per_frame = this->input_stream_info_.frames_to_bytes(1);
  if ((bytes_per_frame == 0) || (bytes_per_frame > RingBufferAudioSource::MAX_ALIGNMENT_BYTES)) {
    return ESP_ERR_NOT_SUPPORTED;
  }

  if ((input_stream_info.get_sample_rate() != output_stream_info.get_sample_rate()) ||
      (input_stream_info.get_bits_per_sample() != output_stream_info.get_bits_per_sample())) {
    this->resampler_ = make_unique<esp_audio_libs::resampler::Resampler>(
        input_stream_info.bytes_to_samples(this->input_buffer_size_),
        output_stream_info.bytes_to_samples(this->output_buffer_size_));

    // Use cascaded biquad filters when downsampling to avoid aliasing
    bool use_pre_filter = output_stream_info.get_sample_rate() < input_stream_info.get_sample_rate();

    esp_audio_libs::resampler::ResamplerConfiguration resample_config = {
        .source_sample_rate = static_cast<float>(input_stream_info.get_sample_rate()),
        .target_sample_rate = static_cast<float>(output_stream_info.get_sample_rate()),
        .source_bits_per_sample = input_stream_info.get_bits_per_sample(),
        .target_bits_per_sample = output_stream_info.get_bits_per_sample(),
        .channels = input_stream_info_.get_channels(),
        .use_pre_or_post_filter = use_pre_filter,
        .subsample_interpolate = false,  // Doubles the CPU load. Using more filters is a better alternative
        .number_of_taps = number_of_taps,
        .number_of_filters = number_of_filters,
    };

    if (!this->resampler_->initialize(resample_config)) {
      // Failed to allocate the resampler's internal buffers
      return ESP_ERR_NO_MEM;
    }
  }

  return ESP_OK;
}

AudioResamplerState AudioResampler::resample(bool stop_gracefully, int32_t *ms_differential) {
  if (this->audio_source_ == nullptr) {
    // Lazily create the zero-copy source on first use. Frame-aligned reads ensure multi-channel frames are
    // never split across the ring buffer's wrap boundary.
    const size_t bytes_per_frame = this->input_stream_info_.frames_to_bytes(1);
    if ((bytes_per_frame == 0) || (bytes_per_frame > RingBufferAudioSource::MAX_ALIGNMENT_BYTES)) {
      // Stream info is unset or the frame is too large to use as an alignment; the uint8_t cast below would
      // truncate it and could yield a source that tears frames.
      return AudioResamplerState::FAILED;
    }
    // Pass the shared_ptr by copy so a failed create() leaves source_ring_buffer_ intact; release our
    // reference only after the source has taken ownership.
    this->audio_source_ = RingBufferAudioSource::create(this->source_ring_buffer_, this->input_buffer_size_,
                                                        static_cast<uint8_t>(bytes_per_frame));
    if (this->audio_source_ == nullptr) {
      return AudioResamplerState::FAILED;
    }
    this->source_ring_buffer_.reset();
  }

  if (stop_gracefully) {
    if (!this->audio_source_->has_buffered_data() && (this->output_transfer_buffer_->available() == 0)) {
      return AudioResamplerState::FINISHED;
    }
  }

  if (!this->pause_output_) {
    // Move audio data to the sink without shifting the data in the output transfer buffer to avoid unnecessary, slow
    // data moves
    this->output_transfer_buffer_->transfer_data_to_sink(pdMS_TO_TICKS(READ_WRITE_TIMEOUT_MS), false);
  } else {
    // If paused, block to avoid wasting CPU resources
    delay(READ_WRITE_TIMEOUT_MS);
  }

  // Expose a chunk of the ring buffer's internal storage. pre_shift is ignored by RingBufferAudioSource
  // (there is no intermediate transfer buffer to compact).
  this->audio_source_->fill(pdMS_TO_TICKS(READ_WRITE_TIMEOUT_MS), false);

  if (this->audio_source_->available() == 0) {
    // No samples available to process
    return AudioResamplerState::RESAMPLING;
  }

  const size_t bytes_free = this->output_transfer_buffer_->free();
  const uint32_t frames_free = this->output_stream_info_.bytes_to_frames(bytes_free);

  const size_t bytes_available = this->audio_source_->available();
  const uint32_t frames_available = this->input_stream_info_.bytes_to_frames(bytes_available);

  if ((this->input_stream_info_.get_sample_rate() != this->output_stream_info_.get_sample_rate()) ||
      (this->input_stream_info_.get_bits_per_sample() != this->output_stream_info_.get_bits_per_sample())) {
    // Adjust gain by -3 dB to avoid clipping due to the resampling process
    esp_audio_libs::resampler::ResamplerResults results =
        this->resampler_->resample(this->audio_source_->data(), this->output_transfer_buffer_->get_buffer_end(),
                                   frames_available, frames_free, -3);

    this->audio_source_->consume(this->input_stream_info_.frames_to_bytes(results.frames_used));
    this->output_transfer_buffer_->increase_buffer_length(
        this->output_stream_info_.frames_to_bytes(results.frames_generated));

    // Resampling causes slight differences in the durations used versus generated. Computes the difference in
    // millisconds. The callback function passing the played audio duration uses the difference to convert from output
    // duration to input duration.
    this->accumulated_frames_used_ += results.frames_used;
    this->accumulated_frames_generated_ += results.frames_generated;

    const int32_t used_ms =
        this->input_stream_info_.frames_to_milliseconds_with_remainder(&this->accumulated_frames_used_);
    const int32_t generated_ms =
        this->output_stream_info_.frames_to_milliseconds_with_remainder(&this->accumulated_frames_generated_);

    *ms_differential = used_ms - generated_ms;

  } else {
    // No resampling required, copy samples directly to the output transfer buffer
    *ms_differential = 0;

    const size_t bytes_to_transfer = std::min(this->output_stream_info_.frames_to_bytes(frames_free),
                                              this->input_stream_info_.frames_to_bytes(frames_available));

    std::memcpy((void *) this->output_transfer_buffer_->get_buffer_end(), (const void *) this->audio_source_->data(),
                bytes_to_transfer);

    this->audio_source_->consume(bytes_to_transfer);
    this->output_transfer_buffer_->increase_buffer_length(bytes_to_transfer);
  }

  return AudioResamplerState::RESAMPLING;
}

}  // namespace esphome::audio

#endif
