#include "transui_mp3.h"

#include "layer3.h"

#include <stdlib.h>
#include <string.h>

static int append_bytes(uint8_t **output, size_t *size, size_t *capacity,
                        const unsigned char *bytes, size_t byte_count) {
  if (byte_count == 0) return TRANSUI_MP3_OK;
  if (*size > SIZE_MAX - byte_count) return TRANSUI_MP3_OUT_OF_MEMORY;
  const size_t required = *size + byte_count;
  if (required > *capacity) {
    size_t next_capacity = *capacity;
    while (next_capacity < required) {
      if (next_capacity > SIZE_MAX / 2) {
        next_capacity = required;
        break;
      }
      next_capacity *= 2;
    }
    uint8_t *resized = realloc(*output, next_capacity);
    if (!resized) return TRANSUI_MP3_OUT_OF_MEMORY;
    *output = resized;
    *capacity = next_capacity;
  }
  memcpy(*output + *size, bytes, byte_count);
  *size = required;
  return TRANSUI_MP3_OK;
}

int transui_mp3_encode_mono_16k_32kbps(const uint8_t *pcm_le,
                                       size_t pcm_byte_count,
                                       uint8_t **output,
                                       size_t *output_size) {
  if (!output || !output_size) return TRANSUI_MP3_INVALID_INPUT;
  *output = NULL;
  *output_size = 0;
  if (!pcm_le || pcm_byte_count < 2 || pcm_byte_count % 2 != 0) {
    return TRANSUI_MP3_INVALID_INPUT;
  }
  if (pcm_byte_count / 8 > SIZE_MAX - 8192) {
    return TRANSUI_MP3_OUT_OF_MEMORY;
  }

  shine_config_t config;
  config.wave.channels = PCM_MONO;
  config.wave.samplerate = 16000;
  shine_set_config_mpeg_defaults(&config.mpeg);
  config.mpeg.mode = MONO;
  config.mpeg.bitr = 32;

  shine_t encoder = shine_initialise(&config);
  if (!encoder) return TRANSUI_MP3_ENCODER_ERROR;

  const int samples_per_pass = shine_samples_per_pass(encoder);
  if (samples_per_pass <= 0) {
    shine_close(encoder);
    return TRANSUI_MP3_ENCODER_ERROR;
  }
  int16_t *frame = calloc((size_t)samples_per_pass, sizeof(int16_t));
  size_t capacity = pcm_byte_count / 8 + 8192;
  uint8_t *encoded = malloc(capacity);
  if (!frame || !encoded) {
    free(frame);
    free(encoded);
    shine_close(encoder);
    return TRANSUI_MP3_OUT_OF_MEMORY;
  }

  const size_t sample_count = pcm_byte_count / 2;
  size_t sample_offset = 0;
  size_t encoded_size = 0;
  int result = TRANSUI_MP3_OK;
  while (sample_offset < sample_count) {
    const size_t remaining = sample_count - sample_offset;
    const size_t frame_samples =
        remaining < (size_t)samples_per_pass ? remaining
                                             : (size_t)samples_per_pass;
    memset(frame, 0, (size_t)samples_per_pass * sizeof(int16_t));
    for (size_t index = 0; index < frame_samples; ++index) {
      const size_t byte_offset = (sample_offset + index) * 2;
      frame[index] = (int16_t)((uint16_t)pcm_le[byte_offset] |
                               ((uint16_t)pcm_le[byte_offset + 1] << 8));
    }

    int written = 0;
    unsigned char *bytes =
        shine_encode_buffer_interleaved(encoder, frame, &written);
    if (written < 0 || (!bytes && written > 0)) {
      result = TRANSUI_MP3_ENCODER_ERROR;
      break;
    }
    result = append_bytes(&encoded, &encoded_size, &capacity, bytes,
                          (size_t)written);
    if (result != TRANSUI_MP3_OK) break;
    sample_offset += frame_samples;
  }

  if (result == TRANSUI_MP3_OK) {
    int written = 0;
    unsigned char *bytes = shine_flush(encoder, &written);
    if (written < 0 || (!bytes && written > 0)) {
      result = TRANSUI_MP3_ENCODER_ERROR;
    } else {
      result = append_bytes(&encoded, &encoded_size, &capacity, bytes,
                            (size_t)written);
    }
  }

  free(frame);
  shine_close(encoder);
  if (result != TRANSUI_MP3_OK || encoded_size == 0) {
    free(encoded);
    return result == TRANSUI_MP3_OK ? TRANSUI_MP3_ENCODER_ERROR : result;
  }

  *output = encoded;
  *output_size = encoded_size;
  return TRANSUI_MP3_OK;
}

void transui_mp3_free(void *memory) { free(memory); }