#include "es8311.h"

#include "drivers/i2s.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

#include "driver/i2s_std.h"

#include "i2c.h"
#include "pins.h"
#include <math.h>
#include <stdint.h>

i2c_master_dev_handle_t es8311_i2c_dev;

int es8311_write(uint8_t reg, uint8_t val) {
  esp_err_t ret = i2c_write_reg(es8311_i2c_dev, reg, val);
  if (ret != ESP_OK) {
    printf("ES8311 write failed: reg=0x%02X, ret=%d\n", reg, ret);
    return -1;
  }
  return 0;
}

static uint8_t db_to_es8311_vol(float db) {
  if (db < -95.5f)
    db = -95.5f;
  if (db > 32.0f)
    db = 32.0f;

  // ES8311 typically uses 0.5 dB steps
  int v = (int)((db + 95.5f) * 2.0f);

  if (v < 0)
    v = 0;
  if (v > 255)
    v = 255;

  return (uint8_t)v;
}

void es8311_set_dac_volume(float vol) {
  es8311_write(ES8311_DAC_REG32, db_to_es8311_vol(vol));
}

void es8311_set_adc_volume(float vol) {
  es8311_write(ES8311_ADC_REG17, db_to_es8311_vol(vol));
}

uint8_t es8311_read(uint8_t reg) {
  uint8_t val;
  esp_err_t ret = i2c_read_reg(es8311_i2c_dev, reg, &val);
  if (ret != ESP_OK) {
    printf("ES8311 read failed: reg=0x%02X, ret=%d\n", reg, ret);
    return 0xFF;
  }
  return val;
}

static void es8311_write_eq_coeff(uint8_t reg_base, uint32_t coeff) {
  coeff &= 0x3FFFFFFF;

  es8311_write(reg_base + 0, (uint8_t)((coeff >> 24) & 0x3F));
  es8311_write(reg_base + 1, (uint8_t)(coeff >> 16));
  es8311_write(reg_base + 2, (uint8_t)(coeff >> 8));
  es8311_write(reg_base + 3, (uint8_t)(coeff));
}
void es8311_set_eq(const es8311_eq_t *eq) {
  if (!eq)
    return;

  // B0
  es8311_write_eq_coeff(0x1D, eq->b0);

  // A1
  es8311_write_eq_coeff(0x21, eq->a1);

  // A2
  es8311_write_eq_coeff(0x25, eq->a2);

  // B1
  es8311_write_eq_coeff(0x29, eq->b1);

  // B2
  es8311_write_eq_coeff(0x2D, eq->b2);
}

void es8311_init(void) {
  printf("ES8311: Starting initialization...\n");

  i2c_device_config_t dev_cfg = {
      .device_address = ES8311_I2C_ADDR,
      .scl_speed_hz = 400000,
  };

  esp_err_t ret = i2c_master_bus_add_device(bus, &dev_cfg, &es8311_i2c_dev);
  if (ret != ESP_OK) {
    printf("ES8311: Failed to add I2C device: %d\n", ret);
    return;
  }

  // Reset sequence
  es8311_write(ES8311_RESET_REG00, 0x1F);
  vTaskDelay(pdMS_TO_TICKS(20));
  es8311_write(ES8311_RESET_REG00, 0x00);
  es8311_write(ES8311_RESET_REG00, 0x80); // Power-on
  vTaskDelay(pdMS_TO_TICKS(20));

  // Start analog circuitry
  es8311_write(ES8311_SYSTEM_REG0D, 0x01);
  vTaskDelay(pdMS_TO_TICKS(20));

  // Configure I2S format - Slave mode
  es8311_write(ES8311_RESET_REG00, 0x80); // Ensure slave mode
  
  es8311_write(ES8311_SDPIN_REG,
               ES8311_SDPIN_FORMAT_I2S | ES8311_SDPIN_WORD_16BIT |
                   ES8311_SDPIN_LR_NORMAL_POLARITY | ES8311_SDPIN_UNMUTE |
                   ES8311_SDPIN_SEL_LEFT_TO_DAC);

  es8311_write(ES8311_SDPOUT_REG,
               ES8311_SDPOUT_FORMAT_I2S | ES8311_SDPOUT_WORD_16BIT |
                   ES8311_SDPOUT_LR_NORMAL_POLARITY | ES8311_SDPOUT_UNMUTE);

                   /*
  // Clock configuration for 48kHz from BCLK
  es8311_write(ES8311_CLK_MANAGER_REG01,
               ES8311_CLK_MANAGER_REG01_MCLK_INV_OFF |
                   ES8311_CLK_MANAGER_REG01_MCLK_ON |
                   ES8311_CLK_MANAGER_REG01_MCLK_SEL_FROM_BCLK |
                   ES8311_CLK_MANAGER_REG01_BCLK_ON);
                   */

                     es8311_write(ES8311_CLK_MANAGER_REG01,
                   ES8311_CLK_MANAGER_REG01_BCLK_ON);

                   /*
  es8311_write(ES8311_CLK_MANAGER_REG02, 0x00);
  es8311_write(ES8311_CLK_MANAGER_REG03, 0x10);
  es8311_write(ES8311_CLK_MANAGER_REG04, 0x10);
  es8311_write(ES8311_CLK_MANAGER_REG05, 0x00);
  */
  
  // CRITICAL: Set LRCK divider
  es8311_write(ES8311_CLK_MANAGER_REG06,
               ES8311_CLK_MANAGER_REG06_CONTINUAL_BCLK |
                   ES8311_CLK_MANAGER_REG06_NORMAL_BCLK |
                   ES8311_CLK_MANAGER_REG06_DIV_BCLK_4);
                   
  es8311_write(ES8311_CLK_MANAGER_REG07,
               ES8311_CLK_MANAGER_REG07_ADCDAT_NORMAL_MODE |
                   ES8311_CLK_MANAGER_REG07_BCLK_LRCLK_NORMAL_MODE);
                  
  // MISSING REGISTER - LRCK divider low byte
  es8311_write(ES8311_CLK_MANAGER_REG08, 0xFF);

  vTaskDelay(pdMS_TO_TICKS(20));

  // Power up and enable everything
  es8311_write(ES8311_SYSTEM_REG0E, ES8311_SYSTEM_REG0E_PGA_ENABLE |
                                        ES8311_SYSTEM_REG0E_PDN_MOD_ENABLE |
                                        ES8311_SYSTEM_REG0E_MOD_NORMAL);

  es8311_write(ES8311_SYSTEM_REG12, 0x00); // DAC on
  es8311_write(ES8311_SYSTEM_REG13, 0x10); // Output enable

  es8311_write(ES8311_SYSTEM_REG0F,
               ES8311_SYSTEM_REG0F_DAC_NORMAL_MODE |
                   ES8311_SYSTEM_REG0F_PGA_NORMAL_MODE |
                   ES8311_SYSTEM_REG0F_PGA_OUTPUT_NORMAL_MODE |
                   ES8311_SYSTEM_REG0F_VCMMOD_NORMAL_MODE |
                   ES8311_SYSTEM_REG0F_ADC_REFERENCE_NORMAL_MODE |
                   ES8311_SYSTEM_REG0F_DAC_REFERENCE_NORMAL_MODE |
                   ES8311_SYSTEM_REG0F_FLASH_NORMAL_MODE |
                   ES8311_SYSTEM_REG0F_INT1_NORMAL_MODE);

  // Microphone configuration
  es8311_write(ES8311_SYSTEM_REG14, ES8311_SYSTEM_REG14_DMIC_OFF |
                                        ES8311_SYSTEM_REG14_MIC1 |
                                        ES8311_SYSTEM_REG14_GAIN_30DB);

  es8311_write(ES8311_ADC_REG16, ES8311_ADC_REG16_ADC_GAIN_24DB);
  
  es8311_write(ES8311_ADC_REG1C, ES8311_ADC_REG1C_HPF_DYNAMIC_HPF |
                                     ES8311_ADC_REG1C_ADCEQ_BYPASS);
                                     
  es8311_write(ES8311_ADC_REG18, ES8311_ADC_REG18_AUTO_LEVEL_CONTROL_ENABLE |
                                     ES8311_ADC_REG18_AUTO_MUTE_DISABLE |
                                     ES8311_ADC_REG18_0_25DB_PER_65536LRCK);

  es8311_write(ES8311_ADC_REG19, ES8311_ADC_REG19_ALC_MIN_LEVEL_NEG_30_1DB |
                                     ES8311_ADC_REG19_ALC_MAX_LEVEL_NEG_6_0DB);

  es8311_write(ES8311_ADC_REG1A, ES8311_ADC_REG1A_AUTOMUTE_NEG_96DB |
                                     ES8311_ADC_REG1A_AUTOMUTE_32768_SAMPLES);

  es8311_write(ES8311_GPIO_REG44, ES8311_GPIO_REG44_ADCDAT_SEL_ADC_ADC |
                                      ES8311_GPIO_REG44_ADCDAT_SEL_NO_LOOPBACK);
                                      //ES8311_GPIO_REG44_ADCDAT_ADC_TO_DAC);

  es8311_write(ES8311_ADC_REG15, ES8311_ADC_REG15_RAMPRATE_0_25_4096LRCK);
  es8311_write(ES8311_DAC_REG37, ES8311_DAC_REG37_RAMPRATE_0_25_4096LRCK);
  
  es8311_set_adc_volume(5);
  es8311_set_dac_volume(5);

  es8311_write(ES8311_DAC_REG31, ES8311_DAC_REG31_MUTE_TO8 |
                                     ES8311_DAC_REG31_DSM_UNMUTE |
                                     ES8311_DAC_REG31_DEM_UNMUTE);

  es8311_write(ES8311_DAC_REG34,
               ES8311_DAC_REG34_DRC_DISABLE |
                   ES8311_DAC_REG34_DRC_WINSIZE_0_25_65536LRCK);

  es8311_write(ES8311_DAC_REG35, ES8311_DAC_REG35_DRC_MAX_LEVEL_NEG_6_0DB |
                                     ES8311_DAC_REG35_DRC_MIN_LEVEL_NEG_30_1DB);
  
  printf("ES8311: Initialization complete\n");
}

int audio_write(const int16_t *samples, size_t count) {
  size_t written;

  esp_err_t r = i2s_channel_write(i2s_tx, samples, count * sizeof(int16_t),
                                  &written, portMAX_DELAY);

  if (r != ESP_OK)
    return -1;

  return written / sizeof(int16_t);
}

int audio_read(int16_t *samples, size_t count) {
  size_t read;

  esp_err_t r = i2s_channel_read(i2s_rx, samples, count * sizeof(int16_t),
                                 &read, portMAX_DELAY);

  if (r != ESP_OK)
    return -1;

  return read / sizeof(int16_t);
}

void audio_beep(void) {
  // ES8311 expects stereo I2S (left + right channels)
  static int16_t buf[ES8311_BEEP_DURATION_SAMPLES * 2]; // *2 for stereo

  for (int i = 0; i < ES8311_BEEP_DURATION_SAMPLES; i++) {
    int phase = (i / ES8311_BEEP_HALF_PERIOD) & 1;
    int16_t sample = phase ? ES8311_BEEP_AMPLITUDE : -ES8311_BEEP_AMPLITUDE;

    buf[i * 2] = sample;     // Left channel
    buf[i * 2 + 1] = sample; // Right channel (same as left for mono)
  }

  int written = audio_write(buf, ES8311_BEEP_DURATION_SAMPLES * 2);
  printf("Beep: wrote %d samples\n", written);
}

void audio_test_tone(void) {
  const int sample_rate = 48000;
  const int duration_ms = 100;
  const int freq_hz = 1000; // A4 note
  const int num_samples = (sample_rate * duration_ms) / 1000;

  int16_t *buf = malloc(num_samples * 2 * sizeof(int16_t));
  if (!buf) {
    printf("Failed to allocate tone buffer\n");
    return;
  }

  for (int i = 0; i < num_samples; i++) {
    float t = (float)i / sample_rate;
    int16_t sample = (int16_t)(sin(2.0f * M_PI * freq_hz * t) * 8000.0f);
    buf[i * 2] = sample;     // Left
    buf[i * 2 + 1] = sample; // Right
  }

  audio_write(buf, num_samples * 2);
  free(buf);
}