#include <iostream>
#include "ns/noise_suppressor.h"
#include "common.h"

#define MIX_INPUT_CHANNELS 2
#define MIX_INPUT_SAMPLES  (10 * MIX_INPUT_SAMPLE_RATE/1000)
#define MIX_INPUT_SAMPLE_RATE 48000

//----------------------------------------------
#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))
// 最大/小音量（db）
#define MIN_DB (-10)
#define MAX_DB (60)
// 最大/小音量： 0: 静音； 100:最大音量
#define MUTE_VOLUME (0)
#define MAX_VOLUME (100)

static int vol_scaler_init(int *scaler, int mindb, int maxdb);
typedef struct VolumeCtlUnit
{
    int scaler[MAX_VOLUME + 1];     // 音量表
    int zeroDb;                     // 0db在scaler中的索引
    // 自定义需要调节的音量
    int micVolume;
    VolumeCtlUnit() {
        // 音量控制器初始化
        zeroDb = vol_scaler_init(scaler, MIN_DB, MAX_DB);
        micVolume = 100;
    }
} volume_ctl_unit_t;
static volume_ctl_unit_t kVolCtrlUnit;

static int vol_scaler_init(int *scaler, int mindb, int maxdb)
{
    double tabdb[MAX_VOLUME + 1];
    double tabf [MAX_VOLUME + 1];
    int    z, i;

    for (i = 0; i < (MAX_VOLUME + 1); i++) {
        // (mindb, maxdb)平均分成(MAX_VOLUME + 1)份
        tabdb[i]  = mindb + (maxdb - mindb) * i / (MAX_VOLUME + 1);
        // dB = 20 * log(A1 / A2)，当A1，A2相等时，db为0
        // 这里以(1 << 14)作为原始声音振幅，得到调节后的振幅(A1),将A1存入音量表中
        tabf [i]  = pow(10.0, tabdb[i] / 20.0);
        scaler[i] = (int)((1 << 14) * tabf[i]); // Q14 fix point
    }

    z = -mindb * (MAX_VOLUME + 1) / (maxdb - mindb);
    z = MAX(z, 0  );
    z = MIN(z, MAX_VOLUME);
    scaler[0] = 0;        // 音量表中，0标识静音
    scaler[z] = (1 << 14);// (mindb, maxdb)的中间值作为0db，即不做增益处理

    return z;
}

static void vol_scaler_run(int16_t *buf, int n, int volume)
{
    /* 简易版
    while (n--) {
        *buf = (*buf) * multiplier / 100.0;
        *buf = std::max((int)*buf, -0x7fff);
        *buf = std::min((int)*buf, 0x7fff);
        buf++;
    }
    */
    int multiplier = kVolCtrlUnit.scaler[volume];
    if (multiplier > (1 << 14)) {
        int32_t v;
        while (n--) {
            v = ((int32_t)*buf * multiplier) >> 14;
            v = MAX(v,-0x7fff);
            v = MIN(v, 0x7fff);
            *buf++ = (int16_t)v;
        }
    } else if (multiplier < (1 << 14)) {
        while (n--) {
            *buf = ((int32_t)*buf * multiplier) >> 14;
            buf++;
        }
    }
}
//----------------------------------------------


// usb声卡设备句柄
static rkStreamPtr usbCaptureDev = nullptr;
static SampleInfo kPcmSampleInfo;

using namespace webrtc;

int main()
{
    // PCM参数
    kPcmSampleInfo.channels       = MIX_INPUT_CHANNELS;
    kPcmSampleInfo.fmt            = SAMPLE_FMT_S16;
    kPcmSampleInfo.sample_rate    = MIX_INPUT_SAMPLE_RATE;
    kPcmSampleInfo.nb_samples     = MIX_INPUT_SAMPLES;

    // 声卡设备初始化
    RkStreamInit(2, capture, kPcmSampleInfo, usbCaptureDev);

    // 采集buffer
    int ret = 0;
    std::shared_ptr<easymedia::SampleBuffer> capData = nullptr;
    int buffer_size = GetSampleSize(kPcmSampleInfo) * kPcmSampleInfo.nb_samples;
    void *ptr = malloc(buffer_size);
    capData = std::make_shared<easymedia::SampleBuffer>(easymedia::MediaBuffer(ptr, buffer_size, -1, ptr, free_memory),
                                                        kPcmSampleInfo);
    assert(capData);

    uint8_t *lineBuffer = nullptr;
    size_t lineSize = buffer_size;
    lineBuffer = (uint8_t*)malloc(lineSize);
    assert(lineBuffer);

    // 降噪初始化
    AudioBuffer audio(MIX_INPUT_SAMPLE_RATE, MIX_INPUT_CHANNELS, 
                        MIX_INPUT_SAMPLE_RATE, MIX_INPUT_CHANNELS, 
                        MIX_INPUT_SAMPLE_RATE, MIX_INPUT_CHANNELS);
    StreamConfig stream_config(MIX_INPUT_SAMPLE_RATE, MIX_INPUT_CHANNELS);
    NsConfig cfg;
    cfg.target_level = NsConfig::SuppressionLevel::k21dB;
    NoiseSuppressor ns(cfg, MIX_INPUT_SAMPLE_RATE, MIX_INPUT_CHANNELS);
    bool split_bands = MIX_INPUT_SAMPLE_RATE > 16000;
    printf("ns config: %d\n", stream_config.num_samples());

    FILE *ns_fp = fopen("/root/ns_out.pcm", "wb");
    FILE *input_fp = fopen("/root/input_out.pcm", "wb");
    while (true)
    {
        // 采集
        size_t read_size = usbCaptureDev->Read(capData->GetPtr(), capData->GetSampleSize(), kPcmSampleInfo.nb_samples);
        if (!read_size && errno != EAGAIN) {
            printf("capture error: %s\n", strerror(errno));
            msleep(10);
            continue ;
        }
        capData->SetSamples(read_size);

        // 分流
        uint8_t *ptr = (uint8_t*)capData->GetPtr();
        for (int i = 0; i < capData->GetSamples(); i++) {
            int size = capData->GetSampleSize() / capData->GetSampleInfo().channels;
            // 左声道 (2.4G)
            // memcpy(wirelessBuffer + i * 2 * size, ptr + i * 2 * size, size);
            // memcpy(wirelessBuffer + (i * 2 + 1) * size, ptr + i * 2 * size, size);
            // 右声道 (LINE IN)
            memcpy(lineBuffer + i * 2 * size, ptr + (i * 2 + 1) * size, size);
            memcpy(lineBuffer + (i * 2 + 1) * size, ptr + (i * 2 + 1) * size, size);
        }

        vol_scaler_run((int16_t *)lineBuffer, capData->GetSamples() * kPcmSampleInfo.channels, kVolCtrlUnit.micVolume);

        if (input_fp) fwrite(lineBuffer, 1, lineSize, input_fp);

        // 降噪
        {
            short *buffer = (short *)lineBuffer;
            audio.CopyFrom(buffer, stream_config);
            if (split_bands)
                audio.SplitIntoFrequencyBands();
            ns.Analyze(audio);
            ns.Process(&audio);
            if (split_bands)
                audio.MergeFrequencyBands();
            audio.CopyTo(stream_config, buffer);
        }
        if (ns_fp) fwrite(lineBuffer, 1, lineSize, ns_fp);
    }
    return 0;
}