/*
 * Copyright (c) Hisilicon Technologies Co., Ltd. 2019-2019. All rights reserved.
 * Description: aenc sample
 * Author: audio
 * Create: 2019-09-17
 */

#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>

#include "uapi_system.h"
#include "uapi_aenc.h"

#include "ha_codec_aacenc.h"
#include "ha_codec_g711.h"
#include "ha_codec_g722.h"
#include "ha_codec_amrnb.h"
#include "ha_codec_amrwb.h"

#include "drv_audio_ext.h"
#include "adp_uapi_ext.h"
#include "securec.h"

typedef struct {
    td_char *sample;
    td_char *in_stream;
    td_char *in_channels;
    td_char *in_sample_rate;
    td_char *out_stream;
    td_char *out_stream_type;
} input_arg;

typedef struct _sample_ctx {
    FILE *in_stream;
    td_u32 in_channels;
    td_u32 in_sample_rate;
    FILE *out_stream;
    td_u32 out_stream_type;

    td_bool stop;
    pthread_t put_thread;
    pthread_t get_thread;

    td_handle aenc;

    td_s32 (*aenc_create)(struct _sample_ctx *aenc_ctx);
} sample_aenc_ctx;

#define AENC_IN_PACKET_SIZE (1024 * 4)
#define AENC_SLEEP_TIME (5 * 1000)       /* 5ms */
#define AENC_SEND_SLEEP_TIME (10 * 1000) /* 10ms */

static td_void audio_frame_init(const sample_aenc_ctx *aenc_ctx, uapi_audio_frame *frame)
{
    td_s32 ret;

    ret = memset_s(frame, sizeof(*frame), 0, sizeof(uapi_audio_frame));
    if (ret != EOK) {
        printf("memset_s fail(0x%x)\n", ret);
        return;
    }

    frame->sample_rate = aenc_ctx->in_sample_rate;
    frame->bit_depth = UAPI_AUDIO_BIT_DEPTH_16;
    frame->channels = aenc_ctx->in_channels;
    frame->pcm_samples = AENC_IN_PACKET_SIZE / (frame->channels * sizeof(td_s16));
    frame->interleaved = TD_TRUE;
}

static td_void *aenc_get_thread(td_void *arg)
{
    td_s32 ret;
    td_u32 frames = 0;
    sample_aenc_ctx *aenc_ctx = (sample_aenc_ctx *)arg;
    uapi_es_buf buf = { 0 };

    while (aenc_ctx->stop == TD_FALSE) {
        ret = uapi_aenc_acquire_stream(aenc_ctx->aenc, &buf, 0);
        if (ret != TD_SUCCESS) {
            usleep(AENC_SLEEP_TIME);
            continue;
        }

        (td_void)fwrite(buf.buf, 1, buf.buf_len, aenc_ctx->out_stream);

        if ((++frames % 0x1000) == 0) { /* print every 1024 frames */
            printf("uapi_aenc_acquire_stream frames = 0x%x\n", frames);
        }

        ret = uapi_aenc_release_stream(aenc_ctx->aenc, &buf);
        if (ret != TD_SUCCESS) {
            printf("uapi_aenc_release_stream failed(0x%x)\n", ret);
            return TD_NULL;
        }
    }

    return TD_NULL;
}

static td_void *aenc_put_thread(td_void *arg)
{
    td_s32 ret;
    td_s32 read_size;
    td_u32 encode_frame = 0;
    td_bool send = TD_TRUE;
    td_void *data = TD_NULL;
    uapi_audio_frame ao_frame;
    sample_aenc_ctx *aenc_ctx = (sample_aenc_ctx *)arg;

    audio_frame_init(aenc_ctx, &ao_frame);

    data = (td_void *)malloc(AENC_IN_PACKET_SIZE);
    if (data == TD_NULL) {
        printf("malloc buffer failed!\n");
        return TD_NULL;
    }

    ao_frame.pcm_buffer = (td_s32 *)data;

    while (aenc_ctx->stop == TD_FALSE) {
        if (send == TD_TRUE) {
            read_size = (td_s32)fread(data, 1, AENC_IN_PACKET_SIZE, aenc_ctx->in_stream);
            if (read_size != AENC_IN_PACKET_SIZE) {
                printf("read file end and rewind!\n");
                rewind(aenc_ctx->in_stream);
                continue;
            }
        }

        ao_frame.frame_index = encode_frame;
        ret = uapi_aenc_send_frame(aenc_ctx->aenc, &ao_frame);
        if (ret == TD_SUCCESS) {
            send = TD_TRUE;
            encode_frame++;
        } else {
            send = TD_FALSE;
            usleep(AENC_SEND_SLEEP_TIME);
        }
    }

    free(data);

    return TD_NULL;
}

static td_s32 aac_aenc_create(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;
    uapi_aenc_attr aenc_attr;
    ha_codec_aac_enc_config private_config;
    ha_codec_enc_param param;
    trans_enc_param_uapi_to_mpi(&(aenc_attr.param), &param);

    if ((aenc_ctx->in_channels != 1) && (aenc_ctx->in_channels != UAPI_AUDIO_CHANNEL_2)) {
        printf("aac encoder only support stero/mono input!\n");
        return TD_FAILURE;
    }

    ret = uapi_aenc_register_encoder("libHA.AUDIO.AAC.encode.so");
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_register_encoder failed(0x%x)\n", ret);
        return ret;
    }

    printf("register \"libHA.AUDIO.AAC.encode.so\" success!\n");

    ret = memset_s(&aenc_attr, sizeof(aenc_attr), 0, sizeof(uapi_aenc_attr));
    if (ret != EOK) {
        printf("memset_s fail(0x%x)\n", ret);
        return ret;
    }

    ret = memset_s(&private_config, sizeof(ha_codec_aac_enc_config),
        0, sizeof(ha_codec_aac_enc_config));
    if (ret != EOK) {
        printf("private_config memset fail(0x%x)\n", ret);
        return ret;
    }

    aenc_attr.aenc_type = UAPI_ACODEC_ID_AAC;
    ha_codec_aac_get_default_config(&private_config);
    private_config.sample_rate = (td_s32)aenc_ctx->in_sample_rate;

    ha_codec_aac_get_enc_default_open_param(&param, (td_void *)&private_config);
    trans_enc_param_mpi_to_uapi(&(aenc_attr.param), &param);
    ret = uapi_aenc_create(&aenc_attr, &aenc_ctx->aenc);
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_create failed(0x%x)\n", ret);
    }

    return ret;
}

static td_s32 g711_aenc_create(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;
    uapi_aenc_attr aenc_attr;
    ha_codec_g711_encode_open_config private_config;
    ha_codec_enc_param param;
    trans_enc_param_uapi_to_mpi(&(aenc_attr.param), &param);

    /* g711 encoder only support 8k/16bit/mono */
    if ((aenc_ctx->in_channels != 1) || (aenc_ctx->in_sample_rate != UAPI_AUDIO_SAMPLE_RATE_8K)) {
        printf("g711 encoder only support 8k/16bit/mono input!\n");
        return TD_FAILURE;
    }

    ret = uapi_aenc_register_encoder("libHA.AUDIO.G711.codec.so");
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_register_encoder failed(0x%x)\n", ret);
        return ret;
    }

    printf("register \"libHA.AUDIO.G711.codec.so\" success!\n");

    ret = memset_s(&aenc_attr, sizeof(aenc_attr), 0, sizeof(uapi_aenc_attr));
    if (ret != EOK) {
        printf("memset_s fail(0x%x)\n", ret);
        return ret;
    }

    ret = memset_s(&private_config, sizeof(ha_codec_g711_encode_open_config),
        0, sizeof(ha_codec_g711_encode_open_config));
    if (ret != EOK) {
        printf("private_config memset fail(0x%x)\n", ret);
        return ret;
    }

    aenc_attr.aenc_type = UAPI_ACODEC_ID_G711;
    private_config.is_alaw = 1;
    private_config.vad = TD_TRUE;

    ha_codec_g711_get_enc_default_open_param(&param, (td_void *)&private_config);
    trans_enc_param_mpi_to_uapi(&(aenc_attr.param), &param);
    aenc_attr.param.sample_per_frame = G711_FRAME_LEN;

    ret = uapi_aenc_create(&aenc_attr, &aenc_ctx->aenc);
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_create failed(0x%x)\n", ret);
    }

    return ret;
}

static td_s32 g722_aenc_create(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;
    uapi_aenc_attr aenc_attr;
    ha_codec_enc_param param;
    trans_enc_param_uapi_to_mpi(&(aenc_attr.param), &param);

    if ((aenc_ctx->in_channels != 1) || (aenc_ctx->in_sample_rate != UAPI_AUDIO_SAMPLE_RATE_16K)) {
        printf("g722 encoder only support 16k/16bit/mono input!\n");
        return TD_FAILURE;
    }

    ret = uapi_aenc_register_encoder("libHA.AUDIO.G722.codec.so");
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_register_encoder failed(0x%x)\n", ret);
        return ret;
    }

    printf("register \"libHA.AUDIO.G722.codec.so\" success!\n");

    ret = memset_s(&aenc_attr, sizeof(aenc_attr), 0, sizeof(uapi_aenc_attr));
    if (ret != EOK) {
        printf("memset_s fail(0x%x)\n", ret);
        return ret;
    }

    aenc_attr.aenc_type = UAPI_ACODEC_ID_G722;

    ha_codec_g722_get_enc_default_open_param(&param);
    trans_enc_param_mpi_to_uapi(&(aenc_attr.param), &param);
    ret = uapi_aenc_create(&aenc_attr, &aenc_ctx->aenc);
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_create failed(0x%x)\n", ret);
    }

    return ret;
}

static td_s32 amr_wb_aenc_create(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;
    uapi_aenc_attr aenc_attr;
    ha_codec_amrwb_encode_open_config private_config;
    ha_codec_enc_param param;
    trans_enc_param_uapi_to_mpi(&(aenc_attr.param), &param);

    if ((aenc_ctx->in_channels != 1) || (aenc_ctx->in_sample_rate != UAPI_AUDIO_SAMPLE_RATE_16K)) {
        printf("amw-wb encoder only support 16k/16bit/mono input!\n");
        return TD_FAILURE;
    }

    ret = uapi_aenc_register_encoder("libHA.AUDIO.AMRWB.codec.so");
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_register_encoder failed(0x%x)\n", ret);
        return ret;
    }

    printf("register \"libHA.AUDIO.AMRWB.codec.so\" success!\n");

    ret = memset_s(&aenc_attr, sizeof(aenc_attr), 0, sizeof(uapi_aenc_attr));
    if (ret != EOK) {
        printf("memset_s fail(0x%x)\n", ret);
        return ret;
    }

    ret = memset_s(&private_config, sizeof(ha_codec_amrwb_encode_open_config),
        0, sizeof(ha_codec_amrwb_encode_open_config));
    if (ret != EOK) {
        printf("private_config memset fail(0x%x)\n", ret);
        return ret;
    }

    aenc_attr.aenc_type = UAPI_ACODEC_ID_AMRWB;
    private_config.format = HA_CODEC_AMRWB_FORMAT_MIME;
    private_config.mode = HA_CODEC_AMRWB_MR2385;
    private_config.dtx = TD_FALSE;

    ha_codec_amrwb_get_enc_default_open_param(&param, (td_void *)&private_config);
    trans_enc_param_mpi_to_uapi(&(aenc_attr.param), &param);
    ret = uapi_aenc_create(&aenc_attr, &aenc_ctx->aenc);
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_create failed(0x%x)\n", ret);
    }

    return ret;
}

static td_s32 amr_nb_aenc_create(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;
    uapi_aenc_attr aenc_attr;
    ha_codec_amrnb_encode_open_config private_config;
    ha_codec_enc_param param;
    trans_enc_param_uapi_to_mpi(&(aenc_attr.param), &param);

    if ((aenc_ctx->in_channels != 1) || (aenc_ctx->in_sample_rate != UAPI_AUDIO_SAMPLE_RATE_8K)) {
        printf("amw-nb encoder only support 8k/16bit/mono input!\n");
        return TD_FAILURE;
    }

    ret = uapi_aenc_register_encoder("libHA.AUDIO.AMRNB.codec.so");
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_register_encoder failed(0x%x)\n", ret);
        return ret;
    }

    printf("register \"libHA.AUDIO.AMRNB.codec.so\" success!\n");

    ret = memset_s(&aenc_attr, sizeof(aenc_attr), 0, sizeof(uapi_aenc_attr));
    if (ret != EOK) {
        printf("memset_s fail(0x%x)\n", ret);
        return ret;
    }

    ret = memset_s(&private_config, sizeof(ha_codec_amrnb_encode_open_config),
        0, sizeof(ha_codec_amrnb_encode_open_config));
    if (ret != EOK) {
        printf("private_config memset fail(0x%x)\n", ret);
        return ret;
    }

    aenc_attr.aenc_type = UAPI_ACODEC_ID_AMRNB;
    private_config.format = HA_CODEC_AMRNB_MIME;
    private_config.mode = HA_CODEC_AMRNB_MR475;
    private_config.dtx = TD_FALSE;

    ha_codec_amrnb_get_enc_default_open_param(&param, (td_void *)&private_config);
    trans_enc_param_mpi_to_uapi(&(aenc_attr.param), &param);
    ret = uapi_aenc_create(&aenc_attr, &aenc_ctx->aenc);
    if (ret != TD_SUCCESS) {
        printf("uapi_aenc_create failed(0x%x)\n", ret);
    }

    return ret;
}

typedef struct {
    const td_char *name;
    td_u32 codec_id;
    td_s32 (*aenc_create)(sample_aenc_ctx *aenc_ctx);
} acodec_map;

static const acodec_map g_aenc_map[] = {
    {"aac", UAPI_ACODEC_ID_AAC, aac_aenc_create},
    {"g711", UAPI_ACODEC_ID_G711, g711_aenc_create},
    {"g722", UAPI_ACODEC_ID_G722, g722_aenc_create},
    {"amrwb", UAPI_ACODEC_ID_AMRWB, amr_wb_aenc_create},
    {"amrnb", UAPI_ACODEC_ID_AMRNB, amr_nb_aenc_create},
};

static td_void aenc_deinit(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;

    ret = uapi_aenc_destroy(aenc_ctx->aenc);
    if (ret != TD_SUCCESS) {
        printf("call uapi_aenc_destroy failed(0x%x)\n", ret);
        return;
    }

    aenc_ctx->aenc = TD_INVALID_HANDLE;

    ret = uapi_aenc_deinit();
    if (ret != TD_SUCCESS) {
        printf("call uapi_aenc_deinit failed.\n");
        return;
    }
}

static const acodec_map *get_aenc_codec_map(const td_char *name)
{
    td_u32 i;
    td_u32 num = sizeof(g_aenc_map) / sizeof(g_aenc_map[0]);

    for (i = 0; i < num; i++) {
        if (strcasecmp(name, g_aenc_map[i].name) == 0) {
            return &g_aenc_map[i];
        }
    }

    return TD_NULL;
}

static td_s32 aenc_init(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;

    if (aenc_ctx->aenc_create == TD_NULL) {
        return TD_FAILURE;
    }

    ret = uapi_aenc_init();
    if (ret != TD_SUCCESS) {
        printf("call uapi_aenc_init failed.\n");
        return ret;
    }

    ret = aenc_ctx->aenc_create(aenc_ctx);
    if (ret != TD_SUCCESS) {
        printf("call aenc_ctx->aenc_create failed.\n");
        (td_void)uapi_aenc_deinit();
        return ret;
    }

    return TD_SUCCESS;
}

static td_void sample_stop(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;

    aenc_ctx->stop = TD_TRUE;
    pthread_join(aenc_ctx->get_thread, TD_NULL);
    pthread_join(aenc_ctx->put_thread, TD_NULL);

    ret = uapi_aenc_stop(aenc_ctx->aenc);
    if (ret != TD_SUCCESS) {
        printf("call uapi_aenc_stop failed(0x%x)\n", ret);
        return;
    }
}

static td_s32 sample_start(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;

    ret = uapi_aenc_start(aenc_ctx->aenc);
    if (ret != TD_SUCCESS) {
        printf("call uapi_aenc_start failed(0x%x)\n", ret);
        return ret;
    }

    aenc_ctx->stop = TD_FALSE;
    pthread_create(&aenc_ctx->get_thread, TD_NULL, aenc_get_thread, aenc_ctx);
    pthread_create(&aenc_ctx->put_thread, TD_NULL, aenc_put_thread, aenc_ctx);

    return TD_SUCCESS;
}

static td_void sample_ctx_deinit(sample_aenc_ctx *aenc_ctx)
{
    if (aenc_ctx->out_stream != TD_NULL) {
        (td_void)fclose(aenc_ctx->out_stream);
        aenc_ctx->out_stream = TD_NULL;
    }

    if (aenc_ctx->in_stream != TD_NULL) {
        (td_void)fclose(aenc_ctx->in_stream);
        aenc_ctx->in_stream = TD_NULL;
    }

    aenc_ctx->aenc = TD_INVALID_HANDLE;
}

static td_s32 sample_ctx_init(const input_arg *arg, sample_aenc_ctx *aenc_ctx)
{
    const acodec_map *map = get_aenc_codec_map(arg->out_stream_type);
    if (map == TD_NULL) {
        return TD_FAILURE;
    }

    aenc_ctx->aenc = TD_INVALID_HANDLE;
    aenc_ctx->aenc_create = map->aenc_create;

    aenc_ctx->in_stream = fopen(arg->in_stream, "rb");
    if (aenc_ctx->in_stream == TD_NULL) {
        printf("open file %s error!\n", arg->in_stream);
        return TD_FAILURE;
    }

    aenc_ctx->in_channels = (td_u32)strtol(arg->in_channels, NULL, 10);       /* 10 is Dec */
    aenc_ctx->in_sample_rate = (td_u32)strtol(arg->in_sample_rate, NULL, 10); /* 10 is Dec */

    aenc_ctx->out_stream = fopen(arg->out_stream, "wb");
    if (aenc_ctx->out_stream == TD_NULL) {
        printf("open file %s error!\n", arg->out_stream);
        goto out;
    }

    /* write magic number of AMR file */
    if (map->codec_id == UAPI_ACODEC_ID_AMRWB) {
        (td_void)fwrite(HA_CODEC_AMRWB_MAGIC_NUMBER, 1,
            strlen(HA_CODEC_AMRWB_MAGIC_NUMBER), aenc_ctx->out_stream);
    } else if (map->codec_id == UAPI_ACODEC_ID_AMRNB) {
        (td_void)fwrite(HA_CODEC_AMRNB_MAGIC_NUMBER, 1,
            strlen(HA_CODEC_AMRNB_MAGIC_NUMBER), aenc_ctx->out_stream);
    }

    return TD_SUCCESS;
out:
    (td_void)fclose(aenc_ctx->in_stream);
    aenc_ctx->in_stream = TD_NULL;
    return TD_FAILURE;
}

static const struct {
    td_s32 (*init)(sample_aenc_ctx *aenc_ctx);
    td_void (*deinit)(sample_aenc_ctx *aenc_ctx);
} g_sample_module[] = {
    {aenc_init, aenc_deinit},
    {sample_start, sample_stop},
};

static td_void sample_deinit(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;
    td_s32 module;
    const td_s32 num = sizeof(g_sample_module) / sizeof(g_sample_module[0]);

    for (module = num - 1; module >= 0; module--) {
        g_sample_module[module].deinit(aenc_ctx);
    }

    ret = uapi_sys_deinit();
    if (ret != TD_SUCCESS) {
        printf("call uapi_sys_deinit failed(0x%x)\n", ret);
        return;
    }
}

static td_s32 sample_init(sample_aenc_ctx *aenc_ctx)
{
    td_s32 ret;
    td_s32 module;
    const td_s32 num = sizeof(g_sample_module) / sizeof(g_sample_module[0]);

    ret = uapi_sys_init();
    if (ret != TD_SUCCESS) {
        printf("call uapi_sys_init failed(0x%x)\n", ret);
        return ret;
    }

    for (module = 0; module < num; module++) {
        ret = g_sample_module[module].init(aenc_ctx);
        if (ret != TD_SUCCESS) {
            printf("call init failed(0x%x)\n", ret);
            goto out;
        }
    }

    return TD_SUCCESS;

out:
    module--;
    for (; module >= 0; module--) {
        g_sample_module[module].deinit(aenc_ctx);
    }
    return ret;
}

static td_void usage(const input_arg *arg)
{
    printf("usage: %s infile_name in_channel in_samplerate outfile_name outfile_type\n"
        "         infile_name: input file name\n"
        "         in_channel:1 2\n"
        "         insamplerate:8000 16000\n"
        "         outfile_name: output file name\n"
        "         outfile_type:aac g711 g722 amrwb amrnb\n",
        arg->sample);
    printf("example: %s ./test.wav 1 8000 ./test.g711 g711\n", arg->sample);
}

td_s32 main(int argc, char *argv[])
{
    td_s32 ret;
    sample_aenc_ctx aenc_ctx;

    input_arg *arg = (input_arg *)argv;

    if (argc != sizeof(input_arg) / sizeof(td_char *)) {
        usage(arg);
        return TD_FAILURE;
    }

    ret = memset_s(&aenc_ctx, sizeof(aenc_ctx), 0, sizeof(sample_aenc_ctx));
    if (ret != EOK) {
        printf("call memset_s failed(0x%x)\n", ret);
        return ret;
    }

    ret = sample_ctx_init(arg, &aenc_ctx);
    if (ret != TD_SUCCESS) {
        printf("call sample_ctx_init failed(0x%x)\n", ret);
        return ret;
    }

    ret = sample_init(&aenc_ctx);
    if (ret != TD_SUCCESS) {
        printf("call sample_init failed(0x%x)\n", ret);
        goto out;
    }

    printf("press any key to exit\n");
    (td_void)getchar();

    sample_deinit(&aenc_ctx);
out:
    sample_ctx_deinit(&aenc_ctx);
    return ret;
}