#include <pm_config.h>
#include <stdbool.h>
#include <stdint.h>
#include <string.h>
#include <zephyr/bluetooth/bluetooth.h>
#include <zephyr/bluetooth/conn.h>
#include <zephyr/bluetooth/gatt.h>
#include <zephyr/bluetooth/l2cap.h>
#include <zephyr/bluetooth/uuid.h>
#include <zephyr/device.h>
#include <zephyr/drivers/gpio.h>
#include <zephyr/drivers/i2s.h>
#include <zephyr/fs/fs.h>
#include <zephyr/fs/littlefs.h>
#include <zephyr/kernel.h>
#include <zephyr/logging/log.h>
#include <zephyr/sys/atomic.h>
#include <zephyr/sys/byteorder.h>
#include <zephyr/sys/ring_buffer.h>
#include <zephyr/sys/util.h>

#include "../drivers/speaker_driver.h"
#include "g711.h"
#include "g727.h"
#include "opus_codec.h"
#include "pcm_recorder.h"

LOG_MODULE_REGISTER(walky_dongle, LOG_LEVEL_INF);
K_MUTEX_DEFINE(i2s_mutex);

typedef enum {
  CODEC_PCM,
  CODEC_G711,
  CODEC_G727,
  CODEC_OPUS,
} audio_codec_t;
audio_codec_t active_Codec = CODEC_PCM;

#define OPUS_SAMPLE_RATE 16000
#define OPUS_FRAME_SAMPLES 320
#define FRAME_SAMPLES 160
#define NARROWBAND_SAMPLE_RATE 8000
#define SAMPLE_RAT 48000
#define AUDIO_PCM_FRAME_SAMPLES FRAME_SAMPLES
#define MAX_PCM_SCRATCH_SAMPLES 512
#define RX_RING_SIZE 4096
#define RX_STAGE_SIZE 2048
#define PCM_QUEUE_FRAMES 16
#define AUDIO_RX_STACK_SIZE 4096
#define AUDIO_RX_PRIORITY 5
#define I2S_WORKER_STACK_SIZE 4096
#define I2S_WORKER_PRIORITY 5
#define BUTTON_DEBOUNCE_MS 300


#define WALKY_SDU_MAX_LEN      512
#define WALKY_RX_MPS           247
#define WALKY_EXTRA_RX_CREDITS 7
/* L2CAP SDU reassembly state (seg_recv path) */
static uint8_t  sdu_buf[WALKY_SDU_MAX_LEN];
static size_t   sdu_buf_len;
static bool     sdu_error;
static int64_t  last_rx_ms = -1;
static uint32_t rx_packet_count;

static struct bt_uuid_128 walky_service_uuid_val = BT_UUID_INIT_128(
    BT_UUID_128_ENCODE(0x497482bd, 0x224d, 0x46cf, 0xacf5, 0xdc3df108c842));

static struct bt_uuid_128 walky_psm_char_uuid_val = BT_UUID_INIT_128(
    BT_UUID_128_ENCODE(0x17120b17, 0xad64, 0x427c, 0x965a, 0x42a183a47624));

#define WALKY_SERVICE_UUID ((const struct bt_uuid *)&walky_service_uuid_val)
#define WALKY_PSM_CHAR_UUID ((const struct bt_uuid *)&walky_psm_char_uuid_val)

static OpusDecoder *opus_dec_state;

struct ad_search_result {
  bool found;
};

NET_BUF_POOL_FIXED_DEFINE(l2cap_rx_pool, 16, BT_L2CAP_SDU_BUF_SIZE(512), CONFIG_BT_CONN_TX_USER_DATA_SIZE, NULL);
K_MSGQ_DEFINE(pcm_queue, sizeof(struct pcm_queue_frame), PCM_QUEUE_FRAMES, 4);
K_THREAD_STACK_DEFINE(audio_rx_stack, AUDIO_RX_STACK_SIZE);
K_THREAD_STACK_DEFINE(i2s_worker_stack, I2S_WORKER_STACK_SIZE);
FS_LITTLEFS_DECLARE_DEFAULT_CONFIG(storage);
K_MSGQ_DEFINE(prompt_queue, MAX_FILENAME_LEN, 5, 4);

static struct bt_conn *default_conn;
static struct bt_gatt_read_params read_params;
static struct bt_l2cap_le_chan l2cap_chan;
static uint16_t discovered_psm;
static struct k_work_delayable gatt_read_work;
static g727_state_t *g727_dec_state;
static uint8_t rx_ring_storage[RX_RING_SIZE];
static struct ring_buf rx_ring;
static uint8_t rx_stage[RX_STAGE_SIZE];
static size_t rx_stage_len;
static struct k_sem audio_rx_sem;
static atomic_t rx_reset_pending;
static int16_t pcm_scratch[MAX_PCM_SCRATCH_SAMPLES];
static struct k_thread audio_rx_thread;
static struct k_thread i2s_thread;
static uint32_t last_btn0_time;
static uint32_t last_btn1_time;
static struct gpio_callback prompt_button0_cb;
static struct gpio_callback prompt_button1_cb;
static struct k_work prompt_work;
static const struct gpio_dt_spec prompt_button0 =
    GPIO_DT_SPEC_GET(DT_ALIAS(sw0), gpios);
static const struct gpio_dt_spec prompt_button1 =
    GPIO_DT_SPEC_GET(DT_ALIAS(sw1), gpios);
static const struct bt_le_conn_param conn_param = BT_LE_CONN_PARAM_INIT(10, 10, 0, 100);

static void start_scan(void);
static void audio_frame_received(const uint8_t *payload, size_t len);
static void audio_pcm_ready(const int16_t *pcm, size_t num_samples);

static void audio_pcm_ready(const int16_t *pcm, size_t num_samples) {
  static uint32_t pcm_count;
  static int64_t last_pcm_ms = -1;

  int64_t now = k_uptime_get();
  int64_t delta = 0;

  if (last_pcm_ms >= 0) {
    delta = now - last_pcm_ms;
  }

  last_pcm_ms = now;
  pcm_count++;
  // LOG_INF("PCM #%u: samples=%zu interval=%lld ms", pcm_count,
  // num_samples,delta);

  size_t offset = 0;
  while (offset < num_samples) {
    size_t chunk = MIN(AUDIO_PCM_FRAME_SAMPLES, num_samples - offset);
    struct pcm_queue_frame frame;
    frame.num_samples = chunk;
    memcpy(frame.samples, &pcm[offset], chunk * sizeof(int16_t));
    if (k_msgq_put(&pcm_queue, &frame, K_NO_WAIT) != 0) {
      LOG_WRN("PCM queue full - dropping %zu samples", chunk);
    }

    offset += chunk;
  }
}

#define BLOCK_SAMPLES 160
extern int16_t i2s_stage[BLOCK_SAMPLES];
extern size_t i2s_stage_len;
int l2cap_recv(struct bt_l2cap_chan *chan, struct net_buf *buf) {
  ARG_UNUSED(chan);

  static int64_t last_rx_ms = -1;
  static uint32_t rx_packet_count = 0;

  size_t len = buf->len;

  if (len < 4) {
    LOG_WRN("BLE packet too short: %zu", len);
    return 0;
  }

  int64_t now_ms = k_uptime_get();
  int64_t delta_ms = 0;

  rx_packet_count++;
  if (last_rx_ms >= 0) {
    delta_ms = now_ms - last_rx_ms;
  }

  last_rx_ms = now_ms;
  uint16_t seq = sys_get_be16(&buf->data[0]);
  uint16_t payload_len = sys_get_be16(&buf->data[2]);

  LOG_INF("T7: APPLICATION RX len=%u", buf->len);
  LOG_INF("BLE RX #%u: SEQ=%u packet=%zu payload=%u interval=%lld ms", rx_packet_count, seq, len, payload_len, delta_ms);
  if ((size_t)(4 + payload_len) > len) {
    LOG_WRN("Invalid frame: SEQ=%u payload=%u packet=%zu", seq, payload_len, len);
    return 0;
  }

  // if (g727_dec_state == NULL) {
  //   LOG_WRN("Decoder state NULL");
  //   return 0;
  // }  

  // if (payload_len > BLOCK_SAMPLES) {
  //   LOG_WRN("Payload too large: %u", payload_len);
  //   return 0;
  // }

  // size_t g711_samples = payload_len;
  // static int16_t decoded_pcm[160];
  // if (g711_samples > ARRAY_SIZE(decoded_pcm)) {
  //   LOG_WRN("G.711 payload too large: %zu", g711_samples);
  //   return 0;
  // }

  // g711_decode_buffer((const int8_t *)&buf->data[4], (int)g711_samples, decoded_pcm);  

  // /* Boost volume by 16x (~24 dB) with saturation clamping */
  // // int16_t *samples = (int16_t *)decoded_pcm;
  // // size_t num_samples = g711_samples / sizeof(int16_t);

  // // for (size_t i = 0; i < num_samples; i++) {
  // //     int32_t amp = (int32_t)samples[i] * 16; /* Adjust multiplier (e.g. 8, 16, 32) */      
  // //     if (amp > 32767) {
  // //         amp = 32767;
  // //     } else if (amp < -32768) {
  // //         amp = -32768;
  // //     }
  // //     samples[i] = (int16_t)amp;
  // // }
  
  // size_t offset = 0;  
  // while (offset < g711_samples) {   
  //   size_t chunk = MIN(AUDIO_PCM_FRAME_SAMPLES, g711_samples - offset);    
  //   // pcm_recorder_write(&decoded_pcm[offset],chunk);

  //   struct pcm_queue_frame frame;
  //   frame.num_samples = chunk;
  //   memcpy(frame.samples, &decoded_pcm[offset], chunk * sizeof(int16_t));
  //   if (k_msgq_put(&pcm_queue, &frame, K_NO_WAIT) != 0) {
  //     LOG_WRN("PCM queue full - dropping %zu samples (SEQ=%u)", chunk, seq);
  //   }
  //   offset += chunk;
  // }

  return 0;
}

static void audio_frame_received(const uint8_t *payload, size_t len) {
  switch (active_Codec) {
    case CODEC_PCM: {
      size_t total_samples = len / 2;
      size_t offset = 0;
      while (offset < total_samples) {
        size_t chunk = MIN(total_samples - offset, ARRAY_SIZE(pcm_scratch));
        for (size_t i = 0; i < chunk; i++) {
          pcm_scratch[i] = sys_get_le16(&payload[(offset + i) * 2]);
        }
        audio_pcm_ready(pcm_scratch, chunk);
        offset += chunk;
      }
      break;
    }

    case CODEC_G711: {
      size_t total_samples = len;
      size_t offset = 0;
      while (offset < total_samples) {
        size_t chunk = MIN(total_samples - offset, ARRAY_SIZE(pcm_scratch));
        g711_decode_buffer((const int8_t *)&payload[offset], (int)chunk,
                           pcm_scratch);
        audio_pcm_ready(pcm_scratch, chunk);
        offset += chunk;
      }

      break;
    }

    case CODEC_G727: {
      if (g727_dec_state != NULL) {
        if (len >= 160) {
          g727_decode(payload, 0, pcm_scratch, 0, 160, g727_dec_state);
          audio_pcm_ready(pcm_scratch, 160);
        }
      }
      break;
    }

    case CODEC_OPUS: {
      if (opus_dec_state != NULL) {
        int32_t decoded =
            opus_codec_decode(opus_dec_state, payload, (int32_t)len,
                              pcm_scratch, OPUS_FRAME_SAMPLES);
        if (decoded > 0) {
          audio_pcm_ready(pcm_scratch, (size_t)decoded);
        }
      }
      break;
    }

    default:
      break;
  }
}

static void prompt_work_handler(struct k_work *work) {
  ARG_UNUSED(work);

  char filename[MAX_FILENAME_LEN];

  while (k_msgq_get(&prompt_queue, &filename, K_NO_WAIT) == 0) {
    /* Guard I2S calls to prevent collision with streaming BLE audio */
    k_mutex_lock(&i2s_mutex, K_FOREVER);
    play_audio(filename);
    k_mutex_unlock(&i2s_mutex);
  }
}

static void prompt_button0_pressed(const struct device *dev,
                                   struct gpio_callback *cb, uint32_t pins) {
  ARG_UNUSED(dev);
  ARG_UNUSED(cb);
  ARG_UNUSED(pins);

  uint32_t now = k_uptime_get_32();
  if ((now - last_btn0_time) < BUTTON_DEBOUNCE_MS) {
    return;
  }
  last_btn0_time = now;

  if (k_msgq_put(&prompt_queue, RECORDED_FILE, K_NO_WAIT) != 0) {
    LOG_WRN("Prompt queue full");
    return;
  }

  k_work_submit(&prompt_work);
}

static void prompt_button1_pressed(const struct device *dev,
                                   struct gpio_callback *cb, uint32_t pins) {
  ARG_UNUSED(dev);
  ARG_UNUSED(cb);
  ARG_UNUSED(pins);

  uint32_t now = k_uptime_get_32();
  if ((now - last_btn1_time) < BUTTON_DEBOUNCE_MS) {
    return;
  }
  last_btn1_time = now;

  if (k_msgq_put(&prompt_queue, ACCESS_DENIED_FILE, K_NO_WAIT) != 0) {
    LOG_WRN("Prompt queue full");
    return;
  }

  k_work_submit(&prompt_work);
}

static void prompt_buttons_init(void) {
  k_work_init(&prompt_work, prompt_work_handler);

  if (gpio_is_ready_dt(&prompt_button0)) {
    gpio_pin_configure_dt(&prompt_button0, GPIO_INPUT);
    gpio_pin_interrupt_configure_dt(&prompt_button0, GPIO_INT_EDGE_TO_ACTIVE);
    gpio_init_callback(&prompt_button0_cb, prompt_button0_pressed,
                       BIT(prompt_button0.pin));
    gpio_add_callback(prompt_button0.port, &prompt_button0_cb);
  }

  if (gpio_is_ready_dt(&prompt_button1)) {
    gpio_pin_configure_dt(&prompt_button1, GPIO_INPUT);
    gpio_pin_interrupt_configure_dt(&prompt_button1, GPIO_INT_EDGE_TO_ACTIVE);
    gpio_init_callback(&prompt_button1_cb, prompt_button1_pressed,
                       BIT(prompt_button1.pin));
    gpio_add_callback(prompt_button1.port, &prompt_button1_cb);
  }
}

void codec_session_start(void) {
  switch (active_Codec) {
    case CODEC_G727:
      if (g727_dec_state == NULL) {
        g727_dec_state = g727_create_state();
        if (g727_dec_state == NULL) {
          LOG_ERR("Failed to create G.727 decoder");
        }
      }
      break;

    case CODEC_OPUS: {
      int err = 0;
      if (opus_dec_state == NULL) {
        opus_dec_state = opus_codec_create_decoder(OPUS_SAMPLE_RATE, &err);
        if (opus_dec_state == NULL) {
          LOG_ERR("Failed to create OPUS decoder err=%d", err);
        }
      }
      break;
    }

    default:
      break;
  }

  uint32_t rate = (active_Codec == CODEC_OPUS) ? OPUS_SAMPLE_RATE : NARROWBAND_SAMPLE_RATE;
  // i2s_set_stream_rate(rate);
  pcm_recorder_start("/extflash/rx_capture.audio");
}

 void codec_session_end(void) {
  if (g727_dec_state != NULL) {
    g727_destroy_state(g727_dec_state);
    g727_dec_state = NULL;
  }

  if (opus_dec_state != NULL) {
    opus_codec_destroy_decoder(opus_dec_state);
    opus_dec_state = NULL;
  }

  k_msgq_purge(&pcm_queue);
  // i2s_set_stream_rate(SAMPLE_RAT);
  pcm_recorder_stop();
}

extern void i2s_worker(void *p1, void *p2, void *p3);
static void audio_pipeline_init(void) {
  ring_buf_init(&rx_ring, sizeof(rx_ring_storage), rx_ring_storage);
  k_sem_init(&audio_rx_sem, 0, 1);  

  k_thread_create(&i2s_thread, i2s_worker_stack,
                  K_THREAD_STACK_SIZEOF(i2s_worker_stack), i2s_worker, NULL,
                  NULL, NULL, I2S_WORKER_PRIORITY, 0, K_NO_WAIT);
  k_thread_name_set(&i2s_thread, "i2s_audio");
  LOG_INF("Audio pipeline initialized");
}

static void l2cap_connected(struct bt_l2cap_chan *chan) {
  ARG_UNUSED(chan);
  struct bt_l2cap_le_chan *le_chan = BT_L2CAP_LE_CHAN(chan);
  LOG_INF("L2CAP channel connected, PSM=0x%04x", discovered_psm);
  LOG_INF("========== L2CAP CONNECTED ==========");
  LOG_INF("RX: CID=0x%04x MTU=%u MPS=%u credits=%d", le_chan->rx.cid, le_chan->rx.mtu, le_chan->rx.mps, atomic_get(&le_chan->rx.credits));
  LOG_INF("TX: CID=0x%04x MTU=%u MPS=%u credits=%d", le_chan->tx.cid, le_chan->tx.mtu, le_chan->tx.mps, atomic_get(&le_chan->tx.credits));  

  i2s_reset_stream();
  codec_session_start();
}

static void l2cap_disconnected(struct bt_l2cap_chan *chan) {
  ARG_UNUSED(chan);

  LOG_INF("L2CAP channel disconnected");
  codec_session_end();
  atomic_set(&rx_reset_pending, 1);
  k_sem_give(&audio_rx_sem);
}

struct net_buf *l2cap_alloc_buf(struct bt_l2cap_chan *chan) {
  ARG_UNUSED(chan);

  return net_buf_alloc(&l2cap_rx_pool, K_NO_WAIT);
}

static const struct bt_l2cap_chan_ops l2cap_ops = {
    .connected = l2cap_connected,
    .disconnected = l2cap_disconnected,
    .recv = l2cap_recv,
    .alloc_buf = l2cap_alloc_buf,    
};

static int open_l2cap_channel(struct bt_conn *conn, uint16_t psm) {
  memset(&l2cap_chan, 0, sizeof(l2cap_chan));
  l2cap_chan.chan.ops = &l2cap_ops;
  return bt_l2cap_chan_connect(conn, &l2cap_chan.chan, psm);
}

static uint8_t psm_read_cb(struct bt_conn *conn, uint8_t err,
                           struct bt_gatt_read_params *params, const void *data,
                           uint16_t length) {
  ARG_UNUSED(conn);
  ARG_UNUSED(params);

  if (err || data == NULL || length < 4) {
    return BT_GATT_ITER_STOP;
  }

  uint32_t psm32 = sys_get_be32((const uint8_t *)data);
  discovered_psm = (uint16_t)psm32;
  LOG_INF("Discovered L2CAP PSM: 0x%04x", discovered_psm);
  open_l2cap_channel(conn, discovered_psm);
  return BT_GATT_ITER_STOP;
}

static void gatt_read_work_handler(struct k_work *work) {
  ARG_UNUSED(work);

  if (!default_conn) {
    return;
  }

  memset(&read_params, 0, sizeof(read_params));
  read_params.func = psm_read_cb;
  read_params.by_uuid.uuid = WALKY_PSM_CHAR_UUID;
  read_params.by_uuid.start_handle = BT_ATT_FIRST_ATTRIBUTE_HANDLE;
  read_params.by_uuid.end_handle = BT_ATT_LAST_ATTRIBUTE_HANDLE;

  int err = bt_gatt_read(default_conn, &read_params);
  if (err) {
    LOG_WRN("bt_gatt_read failed: %d", err);
  }
}

static void connected(struct bt_conn *conn, uint8_t err) {
  if (err) {
    LOG_WRN("BLE connection failed: %u", err);
    if (default_conn) {
      bt_conn_unref(default_conn);
      default_conn = NULL;
    }
    start_scan();
    return;
  }  

  struct bt_conn_info info;
  if (bt_conn_get_info(conn, &info) == 0 && info.type == BT_CONN_TYPE_LE) {
    LOG_INF("Initial conn params: interval=%u (%.2fms) latency=%u timeout=%ums",
            info.le.interval, info.le.interval * 1.25, info.le.latency,
            info.le.timeout * 10);
  }

  LOG_INF("BLE connected");
  k_work_schedule(&gatt_read_work, K_MSEC(100));
}

static void disconnected(struct bt_conn *conn, uint8_t reason) {
  ARG_UNUSED(conn);

  LOG_INF("BLE disconnected, reason=0x%02x", reason);
  k_work_cancel_delayable(&gatt_read_work);
  if (default_conn) {
    bt_conn_unref(default_conn);
    default_conn = NULL;
  }

  codec_session_end();
  atomic_set(&rx_reset_pending, 1);
  k_sem_give(&audio_rx_sem);
  start_scan();
}

BT_CONN_CB_DEFINE(conn_callbacks) = {
    .connected = connected,
    .disconnected = disconnected,
};

static bool ad_has_walky_service(struct bt_data *data, void *user_data) {
  struct ad_search_result *result = user_data;

  if (data->type == BT_DATA_UUID128_ALL || data->type == BT_DATA_UUID128_SOME) {
    for (size_t i = 0; i + 15 < data->data_len; i += 16) {
      if (memcmp(&data->data[i], walky_service_uuid_val.val, 16) == 0) {
        result->found = true;
        return false;
      }
    }
  }

  return true;
}

static void device_found(const bt_addr_le_t *addr, int8_t rssi, uint8_t type,
                         struct net_buf_simple *ad) {
  ARG_UNUSED(type);
  ARG_UNUSED(rssi);

  struct ad_search_result result = {.found = false};
  if (default_conn) {
    return;
  }

  bt_data_parse(ad, ad_has_walky_service, &result);
  if (!result.found) {
    return;
  }

  LOG_INF("Walky device found, connecting");
  bt_le_scan_stop();

  struct bt_conn_le_create_param create_param = BT_CONN_LE_CREATE_PARAM_INIT(
      BT_CONN_LE_OPT_NONE, BT_GAP_SCAN_FAST_INTERVAL, BT_GAP_SCAN_FAST_WINDOW);

  int err = bt_conn_le_create(addr, &create_param, &conn_param, &default_conn);
  if (err) {
    LOG_WRN("bt_conn_le_create failed: %d", err);
    start_scan();
  }
}

static void start_scan(void) {
  struct bt_le_scan_param scan_param = {
      .type = BT_LE_SCAN_TYPE_ACTIVE,
      .options = BT_LE_SCAN_OPT_NONE,
      .interval = 0x0060,
      .window = 0x0030,
  };

  int err = bt_le_scan_start(&scan_param, device_found);
  if (err) {
    LOG_WRN("bt_le_scan_start failed: %d", err);
  }
}

static struct fs_mount_t lfs_mnt = {
    .type = FS_LITTLEFS,
    .fs_data = &storage,
    .storage_dev = (void *)PM_LITTLEFS_STORAGE_ID,
    .mnt_point = "/extflash",
};

int get_pcm_queue_data(struct pcm_queue_frame *frame, k_timeout_t timeout) {
    if (frame == NULL) {
        return -EINVAL;
    }
    return k_msgq_get(&pcm_queue, frame, timeout);
}

int main(void) {
  LOG_INF("BLE Walky Dongle starting");
  active_Codec = CODEC_G727;

  k_work_init_delayable(&gatt_read_work, gatt_read_work_handler);
  audio_pipeline_init();

  if (init_speaker() != 0) {
    LOG_WRN("Continuing without speaker output");
  }

  int fs_err = fs_mount(&lfs_mnt);
  if (fs_err < 0) {
    LOG_WRN("LittleFS mount failed (%d)", fs_err);
  } else {
    LOG_INF("LittleFS mounted at %s", lfs_mnt.mnt_point);
  }

  prompt_buttons_init();

  printf("Starting Bluetooth...\n");  
  int err = bt_enable(NULL);
  if (err) {
    LOG_ERR("Bluetooth init failed (%d)", err);
    return 0;
  }

  LOG_INF("Bluetooth initialized");
  LOG_INF("active_Codec: [%d]\n", active_Codec);
  start_scan();

  return 0;
}