#include <zephyr/kernel.h>
#include <zephyr/device.h>
#include <zephyr/drivers/gpio.h>
#include <zephyr/audio/dmic.h>
#include <zephyr/sys/util.h>
#include <zephyr/logging/log.h>

#include <zephyr/bluetooth/bluetooth.h>
#include <zephyr/bluetooth/gatt.h>
#include <zephyr/bluetooth/uuid.h>
#include <zephyr/bluetooth/services/nus.h>

#ifdef CONFIG_BOARD_XIAO_NRF54LM20A_NRF54LM20A_CPUAPP
#include <zephyr/drivers/regulator.h>
#endif

LOG_MODULE_REGISTER(mic_ble, LOG_LEVEL_INF);

/* ===== Audio parameters ===== */
#define RECORD_TIME_S    10
#define SAMPLE_RATE_HZ   16000
#define SAMPLE_BIT_WIDTH 16
#define BYTES_PER_SAMPLE (SAMPLE_BIT_WIDTH / 8)
#define CHUNK_DURATION_MS 100
#define CHUNK_SIZE_BYTES  (BYTES_PER_SAMPLE * (SAMPLE_RATE_HZ * CHUNK_DURATION_MS) / 1000)
#define CHUNK_COUNT       8
#define TOTAL_CHUNKS      (RECORD_TIME_S * 1000 / CHUNK_DURATION_MS)
#define TOTAL_AUDIO_BYTES (BYTES_PER_SAMPLE * SAMPLE_RATE_HZ * RECORD_TIME_S)
#define READ_TIMEOUT_MS   1000

/* WAV header (44 bytes) + audio data */
#define WAV_HEADER_SIZE   44
#define TOTAL_BUFFER_SIZE (WAV_HEADER_SIZE + TOTAL_AUDIO_BYTES)

/* BLE transfer: max payload per notification (MTU 247 - 3 byte ATT header) */
#define BLE_CHUNK_SIZE    244

/* ===== Application states ===== */
enum app_state {
	STATE_IDLE,
	STATE_RECORDING,
	STATE_BLE_TRANSFERRING,
};

/* ===== Device handles ===== */
static const struct device *const dmic_dev = DEVICE_DT_GET(DT_ALIAS(dmic20));
static const struct gpio_dt_spec green_led =
	GPIO_DT_SPEC_GET(DT_NODELABEL(led2), gpios);
static const struct gpio_dt_spec button =
	GPIO_DT_SPEC_GET(DT_ALIAS(sw0), gpios);

#ifdef CONFIG_BOARD_XIAO_NRF54LM20A_NRF54LM20A_CPUAPP
static const struct device *const power_en_dev =
	DEVICE_DT_GET(DT_NODELABEL(power_en));
static const struct device *const dmic_vdd_dev =
	DEVICE_DT_GET(DT_NODELABEL(dmic_vdd));
#endif

/* ===== Synchronization primitives ===== */
static K_SEM_DEFINE(button_sem, 0, 1);

/* ===== Audio buffer ===== */
static uint8_t audio_buffer[TOTAL_BUFFER_SIZE];
static size_t audio_data_len;

/* ===== Application state ===== */
static enum app_state current_state = STATE_IDLE;
static struct bt_conn *default_conn;
static bool nus_notif_enabled;

/* ===== WAV header template (44 bytes) ===== */
static const uint8_t wav_header_template[WAV_HEADER_SIZE] = {
	'R', 'I', 'F', 'F',
	0, 0, 0, 0,
	'W', 'A', 'V', 'E',
	'f', 'm', 't', ' ',
	16, 0, 0, 0,
	1, 0,
	1, 0,
	0, 0, 0, 0,
	0, 0, 0, 0,
	0, 0,
	0, 0,
	'd', 'a', 't', 'a',
	0, 0, 0, 0,
};

/* ===== DMIC memory pool and config ===== */
K_MEM_SLAB_DEFINE_STATIC(mem_slab, CHUNK_SIZE_BYTES, CHUNK_COUNT, 4);

static struct pcm_stream_cfg stream_cfg = {
	.pcm_rate = SAMPLE_RATE_HZ,
	.pcm_width = SAMPLE_BIT_WIDTH,
	.block_size = CHUNK_SIZE_BYTES,
	.mem_slab = &mem_slab,
};

static struct dmic_cfg dmic_config = {
	.io = {
		.min_pdm_clk_freq = 1000000,
		.max_pdm_clk_freq = 3500000,
		.min_pdm_clk_dc = 40,
		.max_pdm_clk_dc = 60,
	},
	.streams = &stream_cfg,
	.channel = {
		.req_num_streams = 1,
		.req_num_chan = 1,
	},
};

/* ===== Timers ===== */
static struct k_timer record_timer;
static struct k_timer led_blink_timer;

/* ===== BLE transfer work (delayable for retry scheduling) ===== */
static struct k_work_delayable ble_send_work;
static size_t ble_send_offset;
static size_t ble_send_total;

/* ===== Forward declarations ===== */
static int start_recording(void);
static int stop_recording(void);
static void start_ble_transfer(void);

/* ===== Button ISR ===== */
static struct gpio_callback button_cb_data;

static void button_pressed(const struct device *dev, struct gpio_callback *cb,
			   uint32_t pins)
{
	k_sem_give(&button_sem);
}

/* ===== Recording timeout handler ===== */
static void record_timer_handler(struct k_timer *timer)
{
	LOG_INF("Recording timeout");
	if (current_state == STATE_RECORDING) {
		(void)stop_recording();
		start_ble_transfer();
	}
}

/* ===== LED blink timer handler ===== */
static void led_blink_timer_handler(struct k_timer *timer)
{
	if (current_state == STATE_BLE_TRANSFERRING) {
		gpio_pin_toggle_dt(&green_led);
	}
}

/* ===== DMIC power enable (nRF54LM20A) ===== */
#ifdef CONFIG_BOARD_XIAO_NRF54LM20A_NRF54LM20A_CPUAPP
static int enable_dmic_power(void)
{
	int ret;

	if (!device_is_ready(power_en_dev)) {
		LOG_ERR("power_en regulator is not ready");
		return -ENODEV;
	}
	if (!device_is_ready(dmic_vdd_dev)) {
		LOG_ERR("dmic_vdd regulator is not ready");
		return -ENODEV;
	}

	ret = regulator_enable(power_en_dev);
	if (ret < 0 && ret != -EALREADY) {
		LOG_ERR("Failed to enable power_en: %d", ret);
		return ret;
	}
	ret = regulator_enable(dmic_vdd_dev);
	if (ret < 0 && ret != -EALREADY) {
		LOG_ERR("Failed to enable dmic_vdd: %d", ret);
		return ret;
	}

	k_sleep(K_MSEC(20));
	return 0;
}
#endif

/* ===== WAV header writer ===== */
static void wav_header_write(uint8_t *buf, uint32_t data_size)
{
	uint32_t file_size = data_size + 36;
	uint16_t block_align = BYTES_PER_SAMPLE * 1;
	uint32_t byte_rate = SAMPLE_RATE_HZ * block_align;

	memcpy(buf, wav_header_template, WAV_HEADER_SIZE);

	buf[4]  = (uint8_t)(file_size);
	buf[5]  = (uint8_t)(file_size >> 8);
	buf[6]  = (uint8_t)(file_size >> 16);
	buf[7]  = (uint8_t)(file_size >> 24);

	buf[24] = (uint8_t)(SAMPLE_RATE_HZ);
	buf[25] = (uint8_t)(SAMPLE_RATE_HZ >> 8);
	buf[26] = (uint8_t)(SAMPLE_RATE_HZ >> 16);
	buf[27] = (uint8_t)(SAMPLE_RATE_HZ >> 24);

	buf[28] = (uint8_t)(byte_rate);
	buf[29] = (uint8_t)(byte_rate >> 8);
	buf[30] = (uint8_t)(byte_rate >> 16);
	buf[31] = (uint8_t)(byte_rate >> 24);

	buf[32] = (uint8_t)(block_align);
	buf[33] = (uint8_t)(block_align >> 8);

	buf[34] = (uint8_t)(SAMPLE_BIT_WIDTH);
	buf[35] = (uint8_t)(SAMPLE_BIT_WIDTH >> 8);

	buf[40] = (uint8_t)(data_size);
	buf[41] = (uint8_t)(data_size >> 8);
	buf[42] = (uint8_t)(data_size >> 16);
	buf[43] = (uint8_t)(data_size >> 24);
}

/* ===== LED helpers ===== */
static void led_on(void)
{
	gpio_pin_set_dt(&green_led, 1);
}

static void led_off(void)
{
	gpio_pin_set_dt(&green_led, 0);
}

/* ===== Recording ===== */
static int start_recording(void)
{
	int ret;

	LOG_INF("Recording started");
	current_state = STATE_RECORDING;
	audio_data_len = 0;
	led_on();

	memset(audio_buffer, 0, TOTAL_BUFFER_SIZE);

	ret = dmic_configure(dmic_dev, &dmic_config);
	if (ret < 0) {
		LOG_ERR("Failed to configure DMIC: %d", ret);
		current_state = STATE_IDLE;
		led_off();
		return ret;
	}

	ret = dmic_trigger(dmic_dev, DMIC_TRIGGER_START);
	if (ret < 0) {
		LOG_ERR("Failed to start DMIC: %d", ret);
		current_state = STATE_IDLE;
		led_off();
		return ret;
	}

	/* Discard first chunk (startup noise) */
	void *discard_buf;
	uint32_t discard_size;
	ret = dmic_read(dmic_dev, 0, &discard_buf, &discard_size, READ_TIMEOUT_MS);
	if (ret < 0) {
		LOG_WRN("Failed to read discard chunk: %d", ret);
	} else {
		k_mem_slab_free(&mem_slab, discard_buf);
	}

	/* Start recording timeout timer */
	k_timer_start(&record_timer, K_SECONDS(RECORD_TIME_S), K_NO_WAIT);

	return 0;
}

static int stop_recording(void)
{
	int ret;

	LOG_INF("Recording stopped");
	k_timer_stop(&record_timer);

	ret = dmic_trigger(dmic_dev, DMIC_TRIGGER_STOP);
	if (ret < 0) {
		LOG_ERR("Failed to stop DMIC: %d", ret);
	}

	current_state = STATE_BLE_TRANSFERRING;
	LOG_INF("Audio data captured: %u bytes", (uint32_t)audio_data_len);

	return 0;
}

/* ===== Capture audio from DMIC — reads chunks until stopped or timeout ===== */
static int capture_audio_data(void)
{
	int ret;
	void *buffer;
	uint32_t size;
	uint8_t *data_ptr = audio_buffer + WAV_HEADER_SIZE;
	size_t offset = 0;

	for (int i = 0; i < TOTAL_CHUNKS; i++) {
		/* Check for button press (non-blocking) */
		if (k_sem_take(&button_sem, K_NO_WAIT) == 0) {
			LOG_INF("Button pressed, stopping recording early");
			audio_data_len = offset;
			return 0;
		}

		ret = dmic_read(dmic_dev, 0, &buffer, &size, READ_TIMEOUT_MS);
		if (ret < 0) {
			LOG_ERR("DMIC read failed: %d", ret);
			audio_data_len = offset;
			return ret;
		}

		if (offset + CHUNK_SIZE_BYTES <= TOTAL_AUDIO_BYTES) {
			memcpy(data_ptr + offset, buffer, CHUNK_SIZE_BYTES);
			offset += CHUNK_SIZE_BYTES;
		} else {
			LOG_WRN("Audio buffer overflow, discarding chunk");
		}

		k_mem_slab_free(&mem_slab, buffer);

		/* Check for timeout (state changed by timer callback) */
		if (current_state != STATE_RECORDING) {
			LOG_INF("Recording stopped by timeout");
			audio_data_len = offset;
			return 0;
		}
	}

	audio_data_len = offset;
	return 0;
}

/* ===== BLE transfer work handler ===== */
static void ble_send_work_handler(struct k_work *work)
{
	int ret;

	if (!default_conn || !nus_notif_enabled) {
		LOG_WRN("BLE not ready, retrying in 500ms");
		k_work_schedule(&ble_send_work, K_MSEC(500));
		return;
	}

	while (ble_send_offset < ble_send_total) {
		size_t remaining = ble_send_total - ble_send_offset;
		size_t chunk = remaining < BLE_CHUNK_SIZE ? remaining : BLE_CHUNK_SIZE;

		ret = bt_nus_send(default_conn,
				  &audio_buffer[ble_send_offset], chunk);
		if (ret == -EAGAIN) {
			/* BLE stack busy, retry soon */
			k_work_schedule(&ble_send_work, K_MSEC(20));
			return;
		} else if (ret < 0) {
			LOG_ERR("BLE send error: %d, retrying", ret);
			k_work_schedule(&ble_send_work, K_MSEC(200));
			return;
		}

		ble_send_offset += chunk;
	}

	/* Transfer complete */
	LOG_INF("BLE transfer complete");
	k_timer_stop(&led_blink_timer);
	led_off();
	current_state = STATE_IDLE;
	nus_notif_enabled = false;
}

static void start_ble_transfer(void)
{
	uint32_t wav_data_size = (uint32_t)audio_data_len;

	/* Write WAV header at start of buffer (PCM data already at +44) */
	wav_header_write(audio_buffer, wav_data_size);

	ble_send_offset = 0;
	ble_send_total = WAV_HEADER_SIZE + wav_data_size;

	LOG_INF("BLE transfer started, %u bytes total", (uint32_t)ble_send_total);

	/* Start LED blinking during transfer (500ms interval) */
	k_timer_start(&led_blink_timer, K_MSEC(500), K_MSEC(500));

	k_work_schedule(&ble_send_work, K_NO_WAIT);
}

/* ===== BLE NUS callbacks ===== */
static void nus_notif_enabled_cb(bool enabled, void *ctx)
{
	nus_notif_enabled = enabled;
	if (enabled) {
		LOG_INF("NUS notifications enabled by peer");
	}
}

static void nus_received_cb(struct bt_conn *conn, const void *data,
			    uint16_t len, void *ctx)
{
	/* We only send data, not receive */
}

static struct bt_nus_cb nus_cb = {
	.notif_enabled = nus_notif_enabled_cb,
	.received = nus_received_cb,
};

/* ===== BLE connection callbacks ===== */
static void connected(struct bt_conn *conn, uint8_t err)
{
	if (err) {
		LOG_ERR("Connection failed (err %u)", err);
		return;
	}
	default_conn = bt_conn_ref(conn);
	LOG_INF("BLE connected");
}

static void disconnected(struct bt_conn *conn, uint8_t reason)
{
	LOG_INF("BLE disconnected (reason %u)", reason);

	if (default_conn) {
		bt_conn_unref(default_conn);
		default_conn = NULL;
	}
	nus_notif_enabled = false;

	if (current_state == STATE_BLE_TRANSFERRING) {
		LOG_WRN("BLE disconnected during transfer");
		k_timer_stop(&led_blink_timer);
		led_off();
		current_state = STATE_IDLE;
	}
}

static struct bt_conn_cb conn_callbacks = {
	.connected = connected,
	.disconnected = disconnected,
};

/* ===== BLE initialization ===== */
static int ble_init(void)
{
	int ret;

	ret = bt_enable(NULL);
	if (ret) {
		LOG_ERR("BLE enable failed: %d", ret);
		return ret;
	}
	LOG_INF("BLE initialized");

	bt_conn_cb_register(&conn_callbacks);

	ret = bt_nus_cb_register(&nus_cb, NULL);
	if (ret) {
		LOG_ERR("NUS callback register failed: %d", ret);
		return ret;
	}

	return 0;
}

/* ===== BLE advertising ===== */
static int start_advertising(void)
{
	const struct bt_data ad[] = {
		BT_DATA_BYTES(BT_DATA_FLAGS, (BT_LE_AD_GENERAL | BT_LE_AD_NO_BREDR)),
		BT_DATA(BT_DATA_NAME_COMPLETE, CONFIG_BT_DEVICE_NAME,
			sizeof(CONFIG_BT_DEVICE_NAME) - 1),
	};

	return bt_le_adv_start(BT_LE_ADV_CONN_FAST_1, ad, ARRAY_SIZE(ad), NULL, 0);
}

/* ===== Main ===== */
int main(void)
{
	int ret;

	if (!device_is_ready(dmic_dev) || !device_is_ready(green_led.port) ||
	    !device_is_ready(button.port)) {
		LOG_ERR("Required device is not ready.");
		return -ENODEV;
	}

#ifdef CONFIG_BOARD_XIAO_NRF54LM20A_NRF54LM20A_CPUAPP
	ret = enable_dmic_power();
	if (ret < 0) {
		return ret;
	}
#endif

	dmic_config.channel.req_chan_map_lo =
		dmic_build_channel_map(0, 0, PDM_CHAN_LEFT);

	ret = gpio_pin_configure_dt(&green_led, GPIO_OUTPUT_INACTIVE);
	if (ret < 0) {
		LOG_ERR("Failed to configure LED: %d", ret);
		return ret;
	}

	ret = gpio_pin_configure_dt(&button, GPIO_INPUT);
	if (ret < 0) {
		LOG_ERR("Failed to configure button: %d", ret);
		return ret;
	}
	ret = gpio_pin_interrupt_configure_dt(&button, GPIO_INT_EDGE_TO_ACTIVE);
	if (ret < 0) {
		LOG_ERR("Failed to configure button interrupt: %d", ret);
		return ret;
	}
	gpio_init_callback(&button_cb_data, button_pressed, BIT(button.pin));
	gpio_add_callback(button.port, &button_cb_data);

	k_timer_init(&record_timer, record_timer_handler, NULL);
	k_timer_init(&led_blink_timer, led_blink_timer_handler, NULL);
	k_work_init_delayable(&ble_send_work, ble_send_work_handler);

	ret = ble_init();
	if (ret < 0) {
		LOG_ERR("BLE init failed: %d", ret);
		return ret;
	}

	ret = start_advertising();
	if (ret) {
		LOG_ERR("Advertising failed to start: %d", ret);
		return ret;
	}

	LOG_INF("XIAO nRF54LM20A BLE Audio Recorder ready");
	LOG_INF("BOOT button: 1st press records, 2nd press stops & transfers");
	LOG_INF("Max recording: %d seconds", RECORD_TIME_S);

	while (1) {
		LOG_INF("Waiting for BOOT button press...");
		k_sem_take(&button_sem, K_FOREVER);

		/* Simple debounce */
		k_sleep(K_MSEC(50));

		switch (current_state) {
		case STATE_IDLE:
			/* First press: start recording */
			ret = start_recording();
			if (ret < 0) {
				LOG_ERR("Failed to start recording");
				continue;
			}

			ret = capture_audio_data();
			if (ret < 0) {
				LOG_ERR("Audio capture failed");
				current_state = STATE_IDLE;
				led_off();
				continue;
			}

			/* If capture_audio_data returned but state is still RECORDING,
			 * it means button was pressed to stop early. stop_recording
			 * and start_ble_transfer are called below. If state is already
			 * BLE_TRANSFERRING (timeout), skip stop_recording.
			 */
			if (current_state == STATE_RECORDING) {
				(void)stop_recording();
			}
			start_ble_transfer();
			break;

		case STATE_RECORDING:
			/* Second press: stop and transfer */
			/* capture_audio_data will detect the semaphore and return */
			/* This case is reached if button was pressed during the brief
			 * window between start_recording() and capture_audio_data() */
			LOG_INF("Button pressed during recording, stopping");
			(void)stop_recording();
			start_ble_transfer();
			break;

		case STATE_BLE_TRANSFERRING:
			/* Ignore button during transfer */
			LOG_WRN("BLE transfer in progress, button ignored");
			break;
		}
	}

	return 0;
}