#ifndef PW_AUDIO_H
#define PW_AUDIO_H

#include <math.h>
#include <pipewire/pipewire.h>
#include <pthread.h>
#include <spa/param/audio/format-utils.h>

// Shared audio state
typedef struct {
  float left;
  float right;
  float peak_left;
  float peak_right;
  int terminate;
  pthread_mutex_t lock;
  struct pw_main_loop *loop;
  struct pw_stream *stream;
} pw_audio_t;

static pw_audio_t g_audio;

static float amplitude_to_db(float amplitude) {
  if (amplitude <= 0.00001f)
    return -60.0f;
  float db = 20.0f * log10f(amplitude);
  if (db < -60.0f)
    db = -60.0f;
  if (db > 0.0f)
    db = 0.0f;
  return db;
}

// Normalize dB (-60..0) to 0..100
static float db_to_level(float db) {
  if (db <= -60.0f)
    return 0.0f;
  return ((db + 60.0f) / 60.0f) * 100.0f;
}

static void on_process(void *userdata) {
  pw_audio_t *data = userdata;
  struct pw_buffer *b;
  struct spa_buffer *buf;
  float *samples;
  uint32_t n_channels, n_samples;

  if (data->terminate) {
    pw_main_loop_quit(data->loop);
    return;
  }

  if ((b = pw_stream_dequeue_buffer(data->stream)) == NULL)
    return;

  buf = b->buffer;
  if ((samples = buf->datas[0].data) == NULL) {
    pw_stream_queue_buffer(data->stream, b);
    return;
  }

  n_channels = 2; // stereo
  n_samples = buf->datas[0].chunk->size / sizeof(float);

  // Calculate RMS (Root Mean Square) instead of just peak values
  // This gives a more accurate analog needle movement and prevents immediate
  // pegging.
  float sum_l = 0.0f, sum_r = 0.0f;

  for (uint32_t n = 0; n + 1 < n_samples; n += n_channels) {
    float l = samples[n];
    float r = samples[n + 1];
    sum_l += l * l;
    sum_r += r * r;
  }

  float rms_l = sqrtf(sum_l / (n_samples / n_channels));
  float rms_r = sqrtf(sum_r / (n_samples / n_channels));

  float db_l = amplitude_to_db(rms_l);
  float db_r = amplitude_to_db(rms_r);
  float level_l = db_to_level(db_l);
  float level_r = db_to_level(db_r);

  pthread_mutex_lock(&data->lock);
  data->left = level_l;
  data->right = level_r;

  // Peak hold
  if (level_l > data->peak_left)
    data->peak_left = level_l;
  if (level_r > data->peak_right)
    data->peak_right = level_r;
  pthread_mutex_unlock(&data->lock);

  pw_stream_queue_buffer(data->stream, b);
}

static void on_stream_param_changed(void *_data, uint32_t id,
                                    const struct spa_pod *param) {
  pw_audio_t *data = _data;
  struct spa_audio_info format;

  if (param == NULL || id != SPA_PARAM_Format)
    return;

  if (spa_format_parse(param, &format.media_type, &format.media_subtype) < 0)
    return;

  if (format.media_type != SPA_MEDIA_TYPE_audio ||
      format.media_subtype != SPA_MEDIA_SUBTYPE_raw)
    return;
}

static const struct pw_stream_events stream_events = {
    PW_VERSION_STREAM_EVENTS,
    .param_changed = on_stream_param_changed,
    .process = on_process,
};

static void *pw_audio_thread(void *arg) {
  pw_audio_t *data = &g_audio;
  const struct spa_pod *params[1];
  uint8_t buffer[1024];
  struct spa_pod_builder b = SPA_POD_BUILDER_INIT(buffer, sizeof(buffer));

  pw_init(0, NULL);

  data->loop = pw_main_loop_new(NULL);
  if (data->loop == NULL) {
    data->terminate = 1;
    return NULL;
  }

  struct pw_properties *props = pw_properties_new(
      PW_KEY_MEDIA_TYPE, "Audio", PW_KEY_CONFIG_NAME, "client-rt.conf",
      PW_KEY_MEDIA_CATEGORY, "Capture", PW_KEY_MEDIA_ROLE, "Music", NULL);

  pw_properties_set(props, PW_KEY_STREAM_CAPTURE_SINK, "true");

  data->stream =
      pw_stream_new_simple(pw_main_loop_get_loop(data->loop), "vumeter-capture",
                           props, &stream_events, data);

  params[0] = spa_format_audio_raw_build(
      &b, SPA_PARAM_EnumFormat,
      &SPA_AUDIO_INFO_RAW_INIT(.format = SPA_AUDIO_FORMAT_F32));

  pw_stream_connect(data->stream, PW_DIRECTION_INPUT, PW_ID_ANY,
                    PW_STREAM_FLAG_AUTOCONNECT | PW_STREAM_FLAG_MAP_BUFFERS |
                        PW_STREAM_FLAG_RT_PROCESS,
                    params, 1);

  pw_main_loop_run(data->loop);

  pw_stream_destroy(data->stream);
  pw_main_loop_destroy(data->loop);
  pw_deinit();

  return NULL;
}

static int pw_audio_start() {
  g_audio.left = 0;
  g_audio.right = 0;
  g_audio.peak_left = 0;
  g_audio.peak_right = 0;
  g_audio.terminate = 0;
  pthread_mutex_init(&g_audio.lock, NULL);

  pthread_t thread;
  return pthread_create(&thread, NULL, pw_audio_thread, NULL);
}

static void pw_audio_get_levels(float *left, float *right) {
  pthread_mutex_lock(&g_audio.lock);
  *left = g_audio.left;
  *right = g_audio.right;
  pthread_mutex_unlock(&g_audio.lock);
}

static void pw_audio_stop() {
  g_audio.terminate = 1;
  if (g_audio.loop) {
    pw_main_loop_quit(g_audio.loop);
  }
}

#endif