#include "configuration.h"
#if defined(ARCH_ESP32)
#include "AudioModule.h"
#include "MeshService.h"
#include "NodeDB.h"
#include "RTC.h"
#include "Router.h"
#include "FSCommon.h"

#include <soc/sens_reg.h>
#include <soc/sens_struct.h>
#include <assert.h>

/*
    AudioModule
        A interface to send raw codec2 audio data over the mesh network. Based on the example code from the ESP32_codec2 project.
        https://github.com/deulis/ESP32_Codec2

        Codec 2 is a low-bitrate speech audio codec (speech coding) 
        that is patent free and open source develop by David Grant Rowe.
        http://www.rowetel.com/ and https://github.com/drowe67/codec2

    Basic Usage:
        1) Enable the module by setting audio.codec2_enabled to 1.
        2a) Set the pins for the I2S interface if you want to use digital audio. Recommended on TLora is I2S_WS 13/I2S_SD 15/I2S_SIN 2/I2S_SCK 14
        2b) Set the pins (audio.mic_pin / audio.amp_pin) for your preferred microphone and amplifier GPIO pins if you want to use analog audio.
           This is rather heavy on the CPU and not recommended.
           On tbeam, best use:
                audio.mic_chan 6 (GPIO 34)
                audio.amp_pin 14
                audio.ptt_pin 39 
        3) Set audio.bitrate to the desired codec2 rate (CODEC2_3200, CODEC2_2400, CODEC2_1600, CODEC2_1400, CODEC2_1300, CODEC2_1200, CODEC2_700, CODEC2_700B)

    KNOWN PROBLEMS
        * Until the module is initilized by the startup sequence, the amp_pin is in a floating
          state. This may produce a bit of "noise".
        * Will not work on NRF and the Linux device targets (yet?).
*/

#define AMIC 6
#define AAMP 14
#define PTT_PIN 39

ButterworthFilter hp_filter(240, 8000, ButterworthFilter::ButterworthFilter::Highpass, 1);

// Use I2S Processor 0
#define I2S_PORT I2S_NUM_0

#define AUDIO_MODULE_RX_BUFFER 128
#define AUDIO_MODULE_MODE ModuleConfig_AudioConfig_Audio_Baud_CODEC2_700

TaskHandle_t codec2HandlerTask;
AudioModule *audioModule;

#ifdef ARCH_ESP32
// ESP32 doesn't use that flag
#define YIELD_FROM_ISR(x) portYIELD_FROM_ISR()
#else
#define YIELD_FROM_ISR(x) portYIELD_FROM_ISR(x)
#endif

portMUX_TYPE timerMux = portMUX_INITIALIZER_UNLOCKED;
hw_timer_t* adcTimer = NULL;

//int16_t 1KHz sine test tone
int16_t Sine1KHz[8] = { -21210 , -30000, -21210, 0 , 21210 , 30000 , 21210, 0 };
int Sine1KHz_index = 0;

int IRAM_ATTR local_adc1_read(int channel) {
    uint16_t adc_value;
#if CONFIG_IDF_TARGET_ESP32S3
    SENS.sar_meas1_ctrl2.sar1_en_pad = (1 << channel); // only one channel is selected
    while (SENS.sar_slave_addr1.meas_status != 0);
    SENS.sar_meas1_ctrl2.meas1_start_sar = 0;
    SENS.sar_meas1_ctrl2.meas1_start_sar = 1;
    while (SENS.sar_meas1_ctrl2.meas1_done_sar == 0);
    adc_value = SENS.sar_meas1_ctrl2.meas1_data_sar;
#else
    SENS.sar_meas_start1.sar1_en_pad = (1 << channel); // only one channel is selected
    while (SENS.sar_slave_addr1.meas_status != 0);
    SENS.sar_meas_start1.meas1_start_sar = 0;
    SENS.sar_meas_start1.meas1_start_sar = 1;
    while (SENS.sar_meas_start1.meas1_done_sar == 0);
    adc_value = SENS.sar_meas_start1.meas1_data_sar;
#endif    
    return adc_value;
}

IRAM_ATTR void am_onTimer()
{
    portENTER_CRITICAL_ISR(&timerMux); //Enter crital code without interruptions
    if ((audioModule->radio_state == RadioState::tx) && (!moduleConfig.audio.i2s_sd)) {
        audioModule->adc_buffer[audioModule->adc_buffer_index++] = (16 * local_adc1_read(audioModule->mic_chan)) - 32768;

        //If you want to test with a 1KHz tone, comment the line above and descomment the three lines below

        // adc_buffer[adc_buffer_index++] = Sine1KHz[Sine1KHz_index++];
        // if (Sine1KHz_index >= 8)
        // Sine1KHz_index = 0;

        if (audioModule->adc_buffer_index == audioModule->adc_buffer_size) {
            audioModule->adc_buffer_index = 0;
            memcpy((void*)audioModule->speech, (void*)audioModule->adc_buffer, 2 * audioModule->adc_buffer_size);
            BaseType_t xHigherPriorityTaskWoken = pdFALSE;
            vTaskNotifyGiveFromISR(codec2HandlerTask, &xHigherPriorityTaskWoken);
            if (xHigherPriorityTaskWoken == pdTRUE)
                YIELD_FROM_ISR(xHigherPriorityTaskWoken);
        }
    } else if ((audioModule->radio_state == RadioState::rx) && (!moduleConfig.audio.i2s_din)) {
        // ESP32-S3 does not have DAC support
#if !defined(CONFIG_IDF_TARGET_ESP32S3)

        //Get a value from audio_fifo and convert it to 0 - 255 to play it in the ADC
        if (audioModule->fifo.get(&audioModule->sample))
            audioModule->rx_raw_audio_value = (uint8_t)((audioModule->sample + 32768) / 256);

        dacWrite(moduleConfig.audio.amp_pin ? moduleConfig.audio.amp_pin : AAMP, audioModule->rx_raw_audio_value);
#endif
    }
    portEXIT_CRITICAL_ISR(&timerMux); // exit critical code
}

void run_codec2(void* parameter)
{
    while (true) {
        uint32_t tcount = ulTaskNotifyTake(pdFALSE, pdMS_TO_TICKS(10000));

        if (tcount != 0) {
            if (audioModule->radio_state == RadioState::tx) {
            
                for (int i = 0; i < audioModule->adc_buffer_size; i++)
                    audioModule->speech[i] = (int16_t)hp_filter.Update((float)audioModule->speech[i]);

                codec2_encode(audioModule->codec2, audioModule->tx_encode_frame + audioModule->tx_encode_frame_index, audioModule->speech);    

                //increment the pointer where the encoded frame must be saved
                audioModule->tx_encode_frame_index += audioModule->encode_codec_size; 

                //If it this is reached we have a ready trasnmission frame
                if (audioModule->tx_encode_frame_index == audioModule->encode_frame_size)
                {
                    //Transmit it
                    DEBUG_MSG("♪♫♪ Sending %d bytes\n", audioModule->encode_frame_size);
                    audioModule->sendPayload();
                    audioModule->tx_encode_frame_index = 0;
                }
            }
            if (audioModule->radio_state == RadioState::rx) {
                //Make a cycle to get each codec2 frame from the received frame
                for (int i = 0; i < audioModule->rx_encode_frame_index; i += audioModule->encode_codec_size)
                {
                    //Decode the codec2 frame
                    codec2_decode(audioModule->codec2, audioModule->output_buffer, audioModule->rx_encode_frame + i);

                    if (moduleConfig.audio.i2s_din) {
                        size_t bytesOut = 0;

                        i2s_write(I2S_PORT, &audioModule->output_buffer, audioModule->adc_buffer_size, &bytesOut, pdMS_TO_TICKS(40));
                    } else {
                        //Put the decoded audio in the fifo
                        for (int j = 0; j < audioModule->adc_buffer_size; j++)
                            audioModule->fifo.put(audioModule->output_buffer[j]);
                    }
                }
            }
        }
    }
}

AudioModule::AudioModule() : SinglePortModule("AudioModule", PortNum_AUDIO_APP), concurrency::OSThread("AudioModule")
{
    if ((moduleConfig.audio.codec2_enabled) && (myRegion->audioPermitted)) {
        fifo.init();
        DEBUG_MSG("♪♫♪ Setting up codec2 in mode %u", (moduleConfig.audio.bitrate ? moduleConfig.audio.bitrate : AUDIO_MODULE_MODE) - 1);
        codec2 = codec2_create((moduleConfig.audio.bitrate ? moduleConfig.audio.bitrate : AUDIO_MODULE_MODE) - 1);
        codec2_set_lpc_post_filter(codec2, 1, 0, 0.8, 0.2);
        encode_codec_size = (codec2_bits_per_frame(codec2) + 7) / 8;
        encode_frame_num = Constants_DATA_PAYLOAD_LEN / encode_codec_size;
        encode_frame_size = encode_frame_num * encode_codec_size; // max 237 bytes
        adc_buffer_size = codec2_samples_per_frame(codec2);
        DEBUG_MSG(" using %d frames of %d bytes for a total payload length of %d bytes\n", encode_frame_num, encode_codec_size, encode_frame_size);
        xTaskCreate(&run_codec2, "codec2_task", 30000, NULL, 5, &codec2HandlerTask);
    } else {
        DEBUG_MSG("♪♫♪ Codec2 disabled (AudioModule %d, Region %s, permitted %d)\n", moduleConfig.audio.codec2_enabled, myRegion->name, myRegion->audioPermitted);
    }
}

int32_t AudioModule::runOnce()
{
    if ((moduleConfig.audio.codec2_enabled) && (myRegion->audioPermitted)) {
        esp_err_t res;
        if (firstTime) {
            // if we have I2S_SD defined, take samples from digital mic. I2S_DIN means digital output to amp.
            if (moduleConfig.audio.i2s_sd || moduleConfig.audio.i2s_din) {
                // Set up I2S Processor configuration. This will produce 16bit samples at 8 kHz instead of 12 from the ADC
                DEBUG_MSG("♪♫♪ Initializing I2S SD: %d DIN: %d WS: %d SCK:%d\n", moduleConfig.audio.i2s_sd, moduleConfig.audio.i2s_din, moduleConfig.audio.i2s_ws, moduleConfig.audio.i2s_sck);
                i2s_config_t i2s_config = {
                    .mode = (i2s_mode_t)(I2S_MODE_MASTER | (moduleConfig.audio.i2s_sd ? I2S_MODE_RX : 0) | (moduleConfig.audio.i2s_din ? I2S_MODE_TX : 0)),
                    .sample_rate = 8000,
                    .bits_per_sample = I2S_BITS_PER_SAMPLE_16BIT,
                    .channel_format = I2S_CHANNEL_FMT_ONLY_LEFT,
                    .communication_format = (i2s_comm_format_t)(I2S_COMM_FORMAT_STAND_I2S),
                    .intr_alloc_flags = 0,
                    .dma_buf_count = 8,
                    .dma_buf_len = adc_buffer_size, // 320 * 2 bytes
                    .use_apll = false,
                    .tx_desc_auto_clear = true,
                    .fixed_mclk = 0
                };
                res = i2s_driver_install(I2S_PORT, &i2s_config, 0, NULL);
                if(res != ESP_OK)
                    DEBUG_MSG("♪♫♪ Failed to install I2S driver: %d\n", res);

                const i2s_pin_config_t pin_config = {
                    .bck_io_num = moduleConfig.audio.i2s_sck,
                    .ws_io_num = moduleConfig.audio.i2s_ws,
                    .data_out_num = moduleConfig.audio.i2s_din ? moduleConfig.audio.i2s_din : I2S_PIN_NO_CHANGE,
                    .data_in_num = moduleConfig.audio.i2s_sd ? moduleConfig.audio.i2s_sd : I2S_PIN_NO_CHANGE
                };
                res = i2s_set_pin(I2S_PORT, &pin_config);
                if(res != ESP_OK)
                    DEBUG_MSG("♪♫♪ Failed to set I2S pin config: %d\n", res);

                res = i2s_start(I2S_PORT);
                if(res != ESP_OK)
                    DEBUG_MSG("♪♫♪ Failed to start I2S: %d\n", res);
            }

            // setup analogue DAC only if we don't use I2S for output. This is not available on ESP32-S3
#if !defined(CONFIG_IDF_TARGET_ESP32S3)
            if (!moduleConfig.audio.i2s_din)
                DEBUG_MSG("♪♫♪ Initializing DAC on Pin %u\n", moduleConfig.audio.amp_pin ? moduleConfig.audio.amp_pin : AAMP);
#else
            if (!moduleConfig.audio.i2s_din) {
                DEBUG_MSG("♪♫♪ ESP32-S3 does not support DAC. Audio Module Disabled.\n");
                return INT32_MAX;
            }
#endif
            
            if (!moduleConfig.audio.i2s_sd) {
                 // Set up ADC if we don't have a digital microphone.
                DEBUG_MSG("♪♫♪ Initializing ADC on Channel %u\n", moduleConfig.audio.mic_chan ? moduleConfig.audio.mic_chan : AMIC);
                mic_chan = moduleConfig.audio.mic_chan ? (adc1_channel_t)(int)moduleConfig.audio.mic_chan : (adc1_channel_t)AMIC;
                adc1_config_width(ADC_WIDTH_12Bit);
                adc1_config_channel_atten(mic_chan, ADC_ATTEN_DB_6);
                adc1_get_raw(mic_chan);
            }

            if ((!moduleConfig.audio.i2s_sd) || (!moduleConfig.audio.i2s_din)) {
                // Start a timer at 8kHz to sample the ADC and play the audio on the DAC, but only if we have analogue audio to process
                uint32_t cpufreq = getCpuFrequencyMhz();
                switch (cpufreq){
                    case 160:
                        adcTimer = timerBegin(3, 1000, true); // 160 MHz / 1000 = 160KHz
                        break;
                    case 240:
                        adcTimer = timerBegin(3, 1500, true); // 240 MHz / 1500 = 160KHz
                        break;
                    case 320:
                        adcTimer = timerBegin(3, 2000, true); // 320 MHz / 2000 = 160KHz
                        break;
                    case 80:
                    default:
                        adcTimer = timerBegin(3, 500, true); // 80 MHz / 500 = 160KHz
                        break;
                }
                DEBUG_MSG("♪♫♪ Timer CPU Frequency: %u MHz\n", cpufreq);
                timerAttachInterrupt(adcTimer, &am_onTimer, false);
                timerAlarmWrite(adcTimer, 20, true); // Interrupts when counter == 20, 8.000 times a second
                timerAlarmEnable(adcTimer);
            }

            radio_state = RadioState::rx;

            // Configure PTT input
            DEBUG_MSG("♪♫♪ Initializing PTT on Pin %u\n", moduleConfig.audio.ptt_pin ? moduleConfig.audio.ptt_pin : PTT_PIN);
            pinMode(moduleConfig.audio.ptt_pin ? moduleConfig.audio.ptt_pin : PTT_PIN, INPUT);

            firstTime = false;
        } else {
            // Check if PTT is pressed. TODO hook that into Onebutton/Interrupt drive.
            if (digitalRead(moduleConfig.audio.ptt_pin ? moduleConfig.audio.ptt_pin : PTT_PIN) == HIGH) {
                if (radio_state == RadioState::rx) {
                    DEBUG_MSG("♪♫♪ PTT pressed, switching to TX\n");
                    radio_state = RadioState::tx;
                }
            } else {
                if (radio_state == RadioState::tx) {
                    if (tx_encode_frame_index > 0) {
                        // Send the incomplete frame
                        DEBUG_MSG("♪♫♪ Sending %d bytes (incomplete)\n", tx_encode_frame_index);
                        sendPayload();
                    }
                    DEBUG_MSG("♪♫♪ PTT released, switching to RX\n");	
                    tx_encode_frame_index = 0;
                    radio_state = RadioState::rx;
                }
            }
            if ((radio_state == RadioState::tx) && moduleConfig.audio.i2s_sd) {
                // Get I2S data from the microphone and place in data buffer
                size_t bytesIn = 0;
                res = i2s_read(I2S_PORT, adc_buffer + adc_buffer_index, adc_buffer_size - adc_buffer_index, &bytesIn, pdMS_TO_TICKS(40)); // wait 40ms for audio to arrive.

                if (res == ESP_OK) {
                    adc_buffer_index += bytesIn;
                    if (adc_buffer_index == adc_buffer_size) {
                        adc_buffer_index = 0;
                        memcpy((void*)speech, (void*)adc_buffer, 2 * adc_buffer_size);
                        // Notify run_codec2 task that the buffer is ready.
                        radio_state = RadioState::tx;
                        BaseType_t xHigherPriorityTaskWoken = pdFALSE;
                        vTaskNotifyGiveFromISR(codec2HandlerTask, &xHigherPriorityTaskWoken);
                        if (xHigherPriorityTaskWoken == pdTRUE)
                            YIELD_FROM_ISR(xHigherPriorityTaskWoken);
                    }
                }
            }
        }
        return 100;
    } else {
        DEBUG_MSG("♪♫♪ Audio Module Disabled\n");
        return INT32_MAX;
    }
    
}

MeshPacket *AudioModule::allocReply()
{
    auto reply = allocDataPacket(); // Allocate a packet for sending
    return reply;
}

void AudioModule::sendPayload(NodeNum dest, bool wantReplies)
{
    MeshPacket *p = allocReply();
    p->to = dest;
    p->decoded.want_response = wantReplies;

    p->want_ack = false; // Audio is shoot&forget. TODO: Is this really suppressing retransmissions?
    p->priority = MeshPacket_Priority_MAX; // Audio is important, because realtime

    p->decoded.payload.size = tx_encode_frame_index;
    memcpy(p->decoded.payload.bytes, tx_encode_frame, p->decoded.payload.size);

    service.sendToMesh(p);
}

ProcessMessage AudioModule::handleReceived(const MeshPacket &mp)
{
    if ((moduleConfig.audio.codec2_enabled) && (myRegion->audioPermitted)) {
        auto &p = mp.decoded;
        if (getFrom(&mp) != nodeDB.getNodeNum()) {
            memcpy(rx_encode_frame, p.payload.bytes, p.payload.size);
            radio_state = RadioState::rx;
            rx_encode_frame_index = p.payload.size;
            // Notify run_codec2 task that the buffer is ready.
            BaseType_t xHigherPriorityTaskWoken = pdFALSE;
            vTaskNotifyGiveFromISR(codec2HandlerTask, &xHigherPriorityTaskWoken);
            if (xHigherPriorityTaskWoken == pdTRUE)
                YIELD_FROM_ISR(xHigherPriorityTaskWoken);
        }
    }

    return ProcessMessage::CONTINUE;
}

#endif