#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <x264.h>

#include "x264_enc.h"
//#include "logger.h"

struct _X264Context {
    x264_t *handle;
    x264_picture_t *pic_in;
    x264_picture_t *pic_out;
    x264_param_t *param;
    int64_t frame_cnt;

    /* export buffer */
    uint8_t *h264_buf;
    int force_idr;
    int need_reconfig;

    /* 确保pic_in申请后才能free */
    uint8_t *in_falg;
};

void x264_context_destroy(X264Context *ctx)
{
    if (!ctx)
        return;

    if (ctx->pic_in && ctx->in_falg)
    {
        ctx->in_falg = 0;
        x264_picture_clean(ctx->pic_in);
    }
    if (ctx->handle)
        x264_encoder_close(ctx->handle);

    free(ctx->pic_in);
    free(ctx->pic_out);
    free(ctx->param);

    if (ctx->h264_buf)
        free(ctx->h264_buf);

    free(ctx);
}

X264Context *x264_context_create()
{
    X264Context *ctx;

    ctx = calloc(1, sizeof(X264Context));
    if (!ctx)
        return NULL;

    ctx->pic_in = malloc(sizeof(x264_picture_t));
    if (!ctx->pic_in)
        goto fail;
    ctx->pic_out = malloc(sizeof(x264_picture_t));
    if (!ctx->pic_out)
        goto fail;
    ctx->param = calloc(sizeof(x264_param_t), 1);
    if (!ctx->param)
        goto fail;

    x264_param_default_preset(ctx->param, "ultrafast", "fastdecode");
    ctx->param->i_csp = X264_CSP_NV12;
    ctx->param->i_log_level = X264_LOG_INFO;

    x264_picture_init(ctx->pic_out);

    ctx->frame_cnt = 0;

    return ctx;

fail:
    if (ctx->pic_in)
        free(ctx->pic_in);
    if (ctx->pic_out)
        free(ctx->pic_out);
    if (ctx->param)
        free(ctx->param);

    //logger_log(LOG_FATAL, "x264 init failed\n");
    free(ctx);
    return NULL;
}

static int x264_apply_params(X264Context *ctx)
{
    x264_param_apply_profile(ctx->param, x264_profile_names[2]);

    if (ctx->handle)
        x264_encoder_reconfig(ctx->handle, ctx->param);
    else
        ctx->handle = x264_encoder_open(ctx->param);

    return 0;
}

static int x264_params(X264Context *ctx, int i_period, int idr_period,
                       int bitrate)
{
    if (!ctx || (i_period != 0 && idr_period % i_period != 0)) {
        //logger_log(LOG_ERROR, "idr_period must be a multiplier of i_period\n");
        return -1;
    }

    /* TODO: use correctly */
    if (ctx->param->rc.i_bitrate != bitrate ||
            ctx->param->i_keyint_max != idr_period ||
            ctx->param->i_keyint_min != i_period) {
        ctx->param->i_keyint_max = idr_period;
        ctx->param->i_keyint_min = i_period;

        ctx->param->i_bframe = 5;
        ctx->param->rc.i_rc_method = X264_RC_ABR;
        ctx->param->rc.i_vbv_max_bitrate = bitrate;
        ctx->param->rc.i_vbv_buffer_size = bitrate;
        ctx->param->rc.i_bitrate = bitrate;

        ctx->need_reconfig = 1;
    }

    return 0;
}

typedef struct {
    int i_peroid;
    int idr_period;
    int bitrate;
} enc_level;

int x264_params_level(X264Context *ctx, uint8_t lv)
{
    static const enc_level levels[] = {
        {
            /* high level */
            .i_peroid = 3,
	    .idr_period = 12,
	    .bitrate = 10000000,
        },
        {
            /* medium level */
            .i_peroid = 3,
	    .idr_period = 12,
	    .bitrate = 9000000,
        },
        {
            /* low level */
            .i_peroid = 3,
            .idr_period = 12,
            .bitrate = 3000000,
        },
    };
#define ARRAY_SIZE(x) (sizeof(x) / sizeof(x[0]))
    if (lv >= ARRAY_SIZE(levels)) {
        //logger_log(LOG_ERROR, "encode level error\n");
        return -1;
    }

    /* x264 bitrate in Kbps not bps */
    return x264_params(ctx, levels[lv].i_peroid, levels[lv].idr_period,
                       levels[lv].bitrate / 1000);
}

int x264_import_yuv(X264Context *ctx, uint8_t *yuv, int width, int height)
{
    int y_size;

    if (!ctx)
        return -1;

    if (ctx->param->i_width != width || ctx->param->i_height != height) {
        /* update size */
        ctx->param->i_width = width;
        ctx->param->i_height = height;

        //if (ctx->pic_in)
            //x264_picture_clean(ctx->pic_in);
        x264_picture_alloc(ctx->pic_in, X264_CSP_NV12, width, height);
        ctx->in_falg = 1;

        if (ctx->h264_buf)
            free(ctx->h264_buf);


        ctx->h264_buf = malloc(width * height * 2);

        if(ctx->h264_buf == NULL)
            return -1;

        ctx->need_reconfig = 1;
    }

    y_size = width * height;
    /* copy nv12 data to pic_in */
    memcpy(ctx->pic_in->img.plane[0], yuv, y_size);
    memcpy(ctx->pic_in->img.plane[1], yuv + y_size, y_size/2);

    return 0;
}

void x264_force_idr(X264Context *ctx)
{
    ctx->frame_cnt = 0;
    ctx->pic_in->i_type = X264_TYPE_IDR;
    ctx->force_idr = 1;
}

int x264_do_encode(X264Context *ctx, uint8_t **buf)
{
    int ret, len = 0;
    x264_nal_t *nals;
    int num;

    if (!ctx || !buf)
        return -1;

    ctx->pic_in->i_pts = ctx->frame_cnt++;

    if (ctx->need_reconfig) {
        x264_apply_params(ctx);
        ctx->need_reconfig = 0;
    }

    /* do encoding stuff */
    ret = x264_encoder_encode(ctx->handle, &nals, &num, ctx->pic_in,
                              ctx->pic_out);
    if (ret < 0) {
        printf("x264 encode error\n");
        return -1;
    }

    /* export h264 stream */
    for (int i = 0; i < num; ++i) {
        memcpy(ctx->h264_buf + len, nals[i].p_payload, nals[i].i_payload);
        len += nals[i].i_payload;
    }

    if (ctx->force_idr) {
        ctx->pic_in->i_type = X264_TYPE_AUTO;
        ctx->force_idr = 0;
    }

    *buf = ctx->h264_buf;
    return len;
}

int kylin_x264_encoder_delayed_frames( X264Context *ctx, uint8_t **buf, int *frames)
{
    int ret, len = 0;
    x264_nal_t *nals;
    int num;

    *frames= x264_encoder_delayed_frames(ctx->handle);
    ret = x264_encoder_encode(ctx->handle, &nals, &num, NULL,
                                  ctx->pic_out);

    for (int i = 0; i < num; ++i) {
        memcpy(ctx->h264_buf + len, nals[i].p_payload, nals[i].i_payload);
        len += nals[i].i_payload;
    }

    *buf = ctx->h264_buf;
    return len;
}
