//*****************************************************************************
//
//! @file am_vos_audio.c
//!
//! @brief audio processing work
//
//*****************************************************************************

//*****************************************************************************
//
// Copyright (c) 2024, Ambiq Micro, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice,
// this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
// contributors may be used to endorse or promote products derived from this
// software without specific prior written permission.
//
// Third party software included in this distribution is subject to the
// additional license terms as defined in the /docs/licenses directory.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// This is part of revision ambiqvos_r4.5-fdfa8cf6a4 of the AmbiqSuite Development Package.
//
//*****************************************************************************

#include "am_vos_sys_config.h"
#include "am_vos_board_setup.h"

#include "am_vos_utils.h"

#include "am_vos_task.h"
#include "am_vos_init.h"
#include "am_vos_spp.h"
#include "am_vos_audio.h"
#include "am_vos_codec.h"

#if configUSE_BLE
#include "am_vos_ble.h"
#endif // configUSE_BLE

#if configUSE_DSPC_TalkTo
#include "am_vos_awe.h"
#endif // configUSE_DSPC_TalkTo

#if configUSE_AMVOS_AMA
#include "am_vos_ama.h"
#endif // configUSE_AMVOS_AMA

#if configUSE_RTT_RECORDER
#include "am_vos_rtt_recorder.h"
#endif // configUSE_RTT_RECORDER

#if configUSE_AMVOS_ATVV
#include "am_vos_atvv.h"
#endif // configUSE_AMVOS_ATVV

#if configUSE_Sensory_THF || configUSE_Fluent
int16_t g_i16WWDInputBuff[WWE_INPUT_FRAME_LENTH_SAMPLES];
#endif // configUSE_Sensory_THF || configUSE_Fluent

//
// small ring buffer used between SPP and WWE
//
static uint8_t pui8WweInBuff[WWE_INPUT_FRAME_LENTH_SAMPLES * SAMPLE_16BIT * 2];

VosAudInfo g_sVosAud =
{
#if configUSE_AMVOS_ATVV
    .ui16SequenceNumber = 0,
#endif

    .sWweInputBuf.pui8Data = pui8WweInBuff,
    .sWweInputBuf.ui32BufferTail_write = 0,
    .sWweInputBuf.ui32BufferHead_read = 0,
    .sWweInputBuf.ui32OverWriting = 0,
    .sWweInputBuf.ui32Capacity = WWE_INPUT_FRAME_LENTH_SAMPLES * SAMPLE_16BIT * 2,
    
#if configUSE_DSPC_TalkTo
    .sSppInputBuf.ui32SppInChNum = AM_SPP_IN_CHANNEL_NUM,
    .sSppInputBuf.ui32SppInDataFormat = AM_SPP_IN_DATA_FORMAT,
    .sSppInputBuf.ui32SppInBlockSize = PCM_FRAME_SIZE_SAMPLES,

    .sSppOutputBuf.ui32SppOutChNum = AM_SPP_OUT_CHANNEL_NUM,
    .sSppOutputBuf.ui32SppOutDataFormat = AM_SPP_OUT_DATA_FORMAT,
    .sSppOutputBuf.ui32SppOutBlockSize = PCM_FRAME_SIZE_SAMPLES,
#endif // configUSE_DSPC_TalkTo

#if configUSE_AMBIQ_VADv3
    .psVadHandle = NULL,
#endif // configUSE_AMBIQ_VADv3

#if configUSE_SPP_AGC || configUSE_SPP_DRC
    .pvSppAgcHandle = NULL,
#endif // configUSE_SPP_AGC || configUSE_SPP_DRC

#if configUSE_SPP_DRC
    .pvSppCompHandle = NULL,
#endif // configUSE_SPP_DRC
};

//*****************************************************************************
//
// Apollo SPP I/O structure configuration
//
//*****************************************************************************
//
// This process only take 1-channel data into WWE
//
void am_vos_stereo_to_mono_proc(int32_t *nLRSample, int16_t *nMonoSample)
{
    int32_t nSample;

    for (nSample = 0; nSample < PCM_FRAME_SIZE_SAMPLES; nSample++)
    {
        // Without voice pre-processing(e.g. Beamforming), using left MIC's data as a default.
        nMonoSample[nSample] = nLRSample[nSample] & 0xFFFF;              // Left channel
        //nMonoSample[nSample] = (nLRSample[nSample]>>16) & 0xFFFF;        // Right channel
    }
}

//-----------------------------------------------------------------------------
// METHOD:  am_vos_audio_handler
// PURPOSE: Pass new samples to SPP / WWD engine for processing
//-----------------------------------------------------------------------------

void am_vos_audio_handler(int32_t *nLRSample)
{   
#if configUSE_DSPC_TalkTo
    int16_t* pin16LeftChPtr = (int16_t*)g_sVosAud.sSppInputBuf.SppInputArray[0];
    //
    // PCM data to left channel and right channel
    //
#if configUSE_1MIC_SCNR_LAYOUT
    memcpy((void *)pin16LeftChPtr, (void *)nLRSample, PCM_FRAME_SIZE_SAMPLES * AM_SPP_IN_DATA_FORMAT);
#else
    int16_t* pin16RightChPtr = (int16_t*)g_sVosAud.sSppInputBuf.SppInputArray[1];

    for(uint32_t ui32ChIdx = 0; ui32ChIdx < g_sVosAud.sSppInputBuf.ui32SppInChNum; ui32ChIdx++)
    {
        for (int nSample = 0; nSample < PCM_FRAME_SIZE_SAMPLES; nSample++)
        {
            if(ui32ChIdx == 0)
            {
                pin16LeftChPtr[nSample] = (nLRSample[nSample]) & 0xFFFF;
            }
            else if(ui32ChIdx == 1)
            {
                pin16RightChPtr[nSample] = (nLRSample[nSample]>>16) & 0xFFFF;
            }
        }
    }
#endif

#if VOS_MEASURE_SPP_MIPS
    uint32_t TickBeforeSpp = am_vos_timer_benchmark_read();
#endif // VOS_MEASURE_SPP_MIPS

    //
    // audio data push into spp
    //
    am_spp_input_push(&g_sVosAud.sSppInputBuf);
    //
    // spp data process
    //
    am_spp_process_handler();
    //
    // spp processed data pop out
    //
    am_spp_output_pop(&g_sVosAud.sSppOutputBuf);

#if VOS_MEASURE_SPP_MIPS
    if((am_vos_timer_benchmark_read() - TickBeforeSpp) > 0)
    {
        g_sVosSys.ui32SppTimeTick += am_vos_timer_benchmark_read() - TickBeforeSpp;
        g_sVosSys.ui32SppCallNum++;
        am_vos_timer_benchmark_clear_and_start();
    }
#endif // VOS_MEASURE_SPP_MIPS

#else // configUSE_DSPC_TalkTo
#if USE_MIC_SINGLE
    memcpy((void *)(g_sVosAud.sSppOutputBuf.SppOutputArray[0]), (void *)nLRSample, PCM_FRAME_SIZE_SAMPLES * AM_SPP_OUT_DATA_FORMAT);
#else
    am_vos_stereo_to_mono_proc(nLRSample, (int16_t *)g_sVosAud.sSppOutputBuf.SppOutputArray[0]);
#endif
#endif // configUSE_DSPC_TalkTo

#if configUSE_Sensory_THF || configUSE_Fluent
    if(g_sVosSys.bWwdEnabled == true)
    {
        am_vos_ring_buffer_push(&g_sVosAud.sWweInputBuf, g_sVosAud.sSppOutputBuf.SppOutputArray[0], PCM_FRAME_SIZE_SAMPLES * AM_SPP_OUT_DATA_FORMAT, false);
        if(am_vos_get_ring_buffer_status(&g_sVosAud.sWweInputBuf) >= WWE_INPUT_FRAME_LENTH_SAMPLES * AM_SPP_OUT_DATA_FORMAT)
        {
            am_vos_ring_buffer_pop(&g_sVosAud.sWweInputBuf, g_i16WWDInputBuff, WWE_INPUT_FRAME_LENTH_SAMPLES * AM_SPP_OUT_DATA_FORMAT);
#if VOS_MEASURE_WWE_MIPS
#if (AM_VOS_BENCHMARK_TIMER_CLK == AM_VOS_BENCH_TIMER_DWT)
            am_vos_bench_dwt_reset();
            am_vos_bench_dwt_start();
#else
            uint32_t TickBeforeWwe = am_vos_timer_benchmark_read();
#endif
#endif // VOS_MEASURE_WWE_MIPS
            am_vos_engine_process(g_i16WWDInputBuff, WWE_INPUT_FRAME_LENTH_SAMPLES);
#if VOS_MEASURE_WWE_MIPS
#if (AM_VOS_BENCHMARK_TIMER_CLK == AM_VOS_BENCH_TIMER_DWT)
            g_sVosSys.ui32WweTimeTick += am_vos_bench_dwt_getcycle();
            am_vos_bench_dwt_stop();
            g_sVosSys.ui32WweCallNum++;
#else
            if((am_vos_timer_benchmark_read() - TickBeforeWwe) > 0)
            {
                g_sVosSys.ui32WweTimeTick += am_vos_timer_benchmark_read() - TickBeforeWwe;
                g_sVosSys.ui32WweCallNum++;
                am_vos_timer_benchmark_clear_and_start();
            }
#endif
#endif // VOS_MEASURE_WWE_MIPS
        }
    }
    else
    {
        // 
        // we do not need to run keyword detection algorithm
        // flush the buffer
        //
        am_vos_flush_ring_buffer(&g_sVosAud.sWweInputBuf);
    }
#endif // configUSE_Sensory_THF || configUSE_Fluent

#if configUSE_RTT_RECORDER && configUSE_RECORD_DSPC_SPP
    //
    // Record the raw PCM data and send over RTT
    //
    if(g_sVosSys.ui8RecordStartFlag == 1)
        am_vos_rtt_record((void*)g_sVosAud.sSppOutputBuf.SppOutputArray[0], PCM_FRAME_SIZE_SAMPLES * AM_SPP_OUT_DATA_FORMAT);
#endif // configUSE_RTT_RECORDER, configUSE_RECORD_DSPC_SPP

#if configUSE_AUDIO_CODEC
    am_audio_buffer_nested_push(AM_AUDIO_BUFFER_MONO, g_sVosAud.sSppOutputBuf.SppOutputArray[0], PCM_FRAME_SIZE_SAMPLES * AM_SPP_OUT_DATA_FORMAT);
#endif // configUSE_AUDIO_CODEC

#if configUSE_AMVOS_ATVV
    if(g_sVosSys.ui8KwdDetectedFlag || g_sVosSys.ui8PushTalkFlag)
    {
        if (am_vos_is_connected())
        {
#if (!configUSE_VS_CMD)
            am_vos_audio_wwd_disable();
#endif
            if (!g_sVosSys.ui8ProvideSpeechFlag)
            {
                am_vos_atvv_start_search_send();
                am_vos_ble_hid_search_send();
                
                g_sVosSys.ui32MicOpenTimeOutCount = 1;
            }
        }
        else
        {
            AM_VOS_LOG_WARNING("---- Not connected to Android TV.----\n");
        }
        g_sVosSys.ui8KwdDetectedFlag = 0;
        g_sVosSys.ui8PushTalkFlag = 0;
    }
#endif

#if configUSE_AMVOS_ATVV && configUSE_VS_CMD
    if(g_sVosSys.ui8ProvideSpeechFlag && g_sVosSys.ui8CmdDetectedFlag)
    {
        AM_VOS_LOG_INFO("[ATVV] BACK key send.\n");
        am_vos_ble_hid_back_send();
        am_vos_atvv_audio_end_send();
        am_vos_audio_reset_flag_and_buffer();
    }
#endif

    if (g_sVosSys.ui8KwdDetectedFlag || g_sVosSys.ui8PushTalkFlag || g_sVosSys.ui8ProvideSpeechFlag)
    {
#if configUSE_AAD
        am_vos_vad_aad_threshold_to_low();
#endif // configUSE_AAD

#if configUSE_DSPC_TalkTo
        am_vos_awe_scnr_off();
#endif // configUSE_DSPC_TalkTo                                             
        if (g_sVosSys.ui32StreamingTimeOutCount < AUDIO_KWD_TIMEOUT_S * USE_PCM_SAMPLE_RATE)
        {
            if(g_sVosSys.ui32StreamingTimeOutCount == 0)
            {
#if configUSE_BLE
                am_vos_burst_mode_enable();
                if(!g_sVosSys.ui8PushTalkFlag)
                {
                    am_vos_audio_buffer_rewind();  // in conversation mode, dial back only for 100ms...
                }

#if configUSE_AMVOS_AMA
                // keyword detected, disable barge-in before Alexa understands us
                if (am_vos_is_connected())
                {
                    am_vos_audio_wwd_disable();

                    if (!g_sVosSys.ui8ProvideSpeechFlag)
                    {
                        AmaSpeechInfo sAmaSpeechInfo = 
                        {
                            .ui8PushTalkFlag = g_sVosSys.ui8PushTalkFlag,
                            .bPreRollEnabled = configUSE_PREROLL,
                            .i32PreRollDelayMs = g_sVosSys.i32PreRollDelayMs,
                            .i32DelayedSample = g_sVosSys.i32DelayedSample,
                            .i32EndingSampleCnt = g_sVosSys.i32EndingSampleCnt
                        };
                        am_vos_ama_start_speech_send(&sAmaSpeechInfo);
                    }
                }
                else
                {
                    AM_VOS_LOG_WARNING("---- Not connected to Alexa App.----\n");
                }
#elif configUSE_AMVOS_ATVV
                // keyword detected, disable barge-in before Alexa understands us
                if (am_vos_is_connected())
                {
                    g_sVosSys.ui32MicOpenTimeOutCount = 0;
                    am_vos_atvv_audio_start_send();
                }
                else
                {
                    AM_VOS_LOG_WARNING("---- Not connected to Android TV.----\n");
                }     
#endif // configUSE_AMVOS_AMA
#endif // configUSE_BLE
            }
#if configUSE_AMVOS_AMA || configUSE_AMVOS_ATVV
            if (am_vos_is_connected())
#endif // configUSE_AMVOS_AMA
            {
#if configUSE_AUDIO_CODEC
                am_vos_task_send(AM_VOS_TASK_AUD_PROCESSING, AM_VOS_TASK_CODEC,
                                        AM_VOS_MESSAGE_LONG, AM_SPP_OUT_DATA_FORMAT * PCM_FRAME_SIZE_SAMPLES, 
                                        NULL);  
#endif // configUSE_AUDIO_CODEC
            }
            g_sVosSys.ui32StreamingTimeOutCount += PCM_FRAME_SIZE_SAMPLES;
        }
        else
        {
#if configUSE_AMVOS_ATVV
            am_vos_atvv_audio_end_send();
#endif // configUSE_AMVOS_ATVV
            am_vos_audio_reset_flag_and_buffer();
        }
    }
}

void am_vos_reset_detected_flag()
{
    g_sVosSys.ui8KwdDetectedFlag = 0;
    g_sVosSys.ui8PushTalkFlag = 0;
    g_sVosSys.ui8ProvideSpeechFlag = 0;
    g_sVosSys.ui32StreamingTimeOutCount = 0;
    g_sVosSys.ui8CmdDetectedFlag = 0;

#if configUSE_AMVOS_ATVV
    g_sVosAud.ui16SequenceNumber = 0;
#endif

    am_vos_audio_wwd_enable();

#if configUSE_DSPC_TalkTo && configUSE_TalkTo_SCNR
    am_vos_awe_scnr_on();
#endif // configUSE_DSPC_TalkTo && configUSE_TalkTo_SCNR

#if !(configUSE_BURST_MODE_ALWAYS_ON || configUSE_AIC_LITE_LAYOUT || configUSE_Fluent)
    // AIC layout requires burst mode always on
    am_vos_burst_mode_disable();
#endif // !(configUSE_BURST_MODE_ALWAYS_ON || configUSE_AIC_LITE_LAYOUT || configUSE_Fluent)
}

void am_vos_audio_flush_ring_buffer(void)
{
    am_audio_buffer_init();
}

void am_vos_audio_buffer_rewind(void)
{
    AM_CRITICAL_BEGIN_VOS;
#if (configUSE_AMVOS_AMA && configUSE_PREROLL)

    uint32_t bytes_to_rewind = 0;
    uint32_t bytes_left_in_buffer = 0;
    bytes_left_in_buffer = am_vos_get_ring_buffer_status(&g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO]);
    bytes_to_rewind = (g_sVosSys.i32DelayedSample * SAMPLE_16BIT + SAMPLE_16BIT *
                       (USE_PCM_SAMPLE_RATE / 1000) * AUDIO_PREROLL_TIME_MS);
    if(bytes_to_rewind > bytes_left_in_buffer)
    {
        bytes_to_rewind = bytes_left_in_buffer;
        if(g_sVosSys.i32DelayedSample * SAMPLE_16BIT >= bytes_left_in_buffer)
        {
            g_sVosSys.i32PreRollDelayMs = 1;      // use minimum preroll value
            g_sVosSys.i32DelayedSample = (bytes_left_in_buffer / SAMPLE_16BIT);
        }
        else
        {
            g_sVosSys.i32PreRollDelayMs = (bytes_left_in_buffer / SAMPLE_16BIT -
                                           g_sVosSys.i32DelayedSample) / (USE_PCM_SAMPLE_RATE / 1000);
            //g_sVosSys.i32PreRollDelayMs = (bytes_left_in_buffer / SAMPLE_16BIT) / (USE_PCM_SAMPLE_RATE/1000);
        }
    }
    else
    {
        g_sVosSys.i32PreRollDelayMs = 500;    // use requested preroll value
    }

    AM_VOS_LOG_DEBUG("=== bytes_left_in_buffer = %d, bytes_to_rewind = %d, \n", bytes_left_in_buffer, bytes_to_rewind);
    //AM_VOS_LOG_DEBUG("=== g_sVosSysStatus.i32PreRollDelayMs = %d ms, g_sVosSysStatus.i32DelayedSample = %d samples, \n", g_sVosSys.i32PreRollDelayMs, g_sVosSys.i32DelayedSample);
    AM_VOS_LOG_DEBUG("=== i32PreRollDelayMs = %d ms, i32DelayedSample = %d samples, \n", g_sVosSys.i32PreRollDelayMs, g_sVosSys.i32DelayedSample);

    AM_VOS_LOG_DEBUG("=== before: mono_read_head = %d, mono_write_tail = %d, \n", g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO].ui32BufferHead_read, g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO].ui32BufferTail_write);
    am_audio_buffer_rewind(AM_AUDIO_BUFFER_MONO, ( bytes_to_rewind / CODEC_IN_RING_BUFF_SIZE + 1) * CODEC_IN_RING_BUFF_SIZE); // start index + 500ms preroll, 32000bytes of samples = 1000ms
    AM_VOS_LOG_DEBUG("=== after: mono_read_head = %d, mono_write_tail = %d, \n", g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO].ui32BufferHead_read, g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO].ui32BufferTail_write);

#elif configUSE_AMVOS_ATVV
    am_audio_buffer_rewind(AM_AUDIO_BUFFER_MONO, (320 / CODEC_IN_RING_BUFF_SIZE + 1) * CODEC_IN_RING_BUFF_SIZE); // rewind 10ms voice to overcome the detection delay
#elif configUSE_Fluent
    AM_VOS_LOG_DEBUG("=== before: mono_read_head = %d, write_tail = %d, \n", g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO].ui32BufferHead_read, g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO].ui32BufferTail_write);
    AM_VOS_LOG_DEBUG("=== before: stereo_read_head = %d, write_tail = %d, \n", g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_STEREO].ui32BufferHead_read, g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_STEREO].ui32BufferTail_write);
    am_audio_buffer_rewind(AM_AUDIO_BUFFER_MONO, (320 / CODEC_IN_RING_BUFF_SIZE + 1) * CODEC_IN_RING_BUFF_SIZE); // rewind 20ms voice to overcome the detection delay
    am_audio_buffer_rewind(AM_AUDIO_BUFFER_STEREO, (1280 / CODEC_IN_RING_BUFF_SIZE + 1) * CODEC_IN_RING_BUFF_SIZE); // rewind 40ms voice to overcome the detection delay
    AM_VOS_LOG_DEBUG("=== after: mono_read_head = %d, write_tail = %d, \n", g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO].ui32BufferHead_read, g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_MONO].ui32BufferTail_write);
    AM_VOS_LOG_DEBUG("=== after: stereo_read_head = %d, write_tail = %d, \n", g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_STEREO].ui32BufferHead_read, g_sAmUtil.sRingBuf[AM_AUDIO_BUFFER_STEREO].ui32BufferTail_write);
#else
    am_audio_buffer_rewind(AM_AUDIO_BUFFER_MONO, (3200 / CODEC_IN_RING_BUFF_SIZE + 1) * CODEC_IN_RING_BUFF_SIZE); // rewind 100ms voice to overcome the detection delay
#endif // configUSE_PREROLL
    AM_CRITICAL_END_VOS;
}

void am_vos_audio_reset_flag_and_buffer(void)
{
    am_vos_audio_flush_ring_buffer();
    am_vos_reset_detected_flag();
}

void am_vos_vad_aad_threshold_to_high(void)
{
  //
  //  Change VAD adn AAD threshold to higher (SNR) value until max.
  //
}

void am_vos_vad_aad_threshold_to_low(void)
{
  //
  //  Change VAD and AAD threshold to lower (SNR) value until min.
  //
}

void am_vos_audio_wwd_enable(void)
{
    if(g_sVosSys.bWwdEnabled == false)
        AM_VOS_LOG_INFO("[AM-VoS] Enable WakeWord Detection\n");
    g_sVosSys.bWwdEnabled = true;
    
    return;
}

void am_vos_audio_wwd_disable(void)
{
    if(g_sVosSys.bWwdEnabled == true)
        AM_VOS_LOG_INFO("[AM-VoS] Disable WakeWord Detection\n");
    g_sVosSys.bWwdEnabled = false;

    return;
}

