#include "mpp.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <linux/videodev2.h>
#include <rockchip/mpp_buffer.h>
#include <rockchip/mpp_err.h>
#include <rockchip/mpp_frame.h>
#include <rockchip/mpp_packet.h>
#include <rockchip/rk_mpi.h>

#define WIDTH 1920
#define HEIGHT 1080
#define NUM_PLANES 1
#define BUFFER_COUNT 4

#define SPS_MAX_SIZE 64

static unsigned char sps_buffer[SPS_MAX_SIZE];
static int sps_size = 0;

static int video_fd = -1;
static void *buffers[BUFFER_COUNT];
static struct v4l2_buffer v4l2_buf;
static MppCtx mpp_ctx = NULL;
static MppApi *mpp_api = NULL;
static MppEncCfg cfg;

int video_init(const char *device) {
    video_fd = open(device, O_RDWR);
    if (video_fd < 0) {
        perror("Failed to open video device");
        return -1;
    }

    // Set video format
    struct v4l2_format fmt;
    memset(&fmt, 0, sizeof(fmt));
    fmt.type = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
    fmt.fmt.pix_mp.width = WIDTH;
    fmt.fmt.pix_mp.height = HEIGHT;
    fmt.fmt.pix_mp.pixelformat = V4L2_PIX_FMT_BGR24;
    fmt.fmt.pix_mp.num_planes = NUM_PLANES;

    if (ioctl(video_fd, VIDIOC_S_FMT, &fmt) < 0) {
        perror("Failed to set video format");
        close(video_fd);
        return -1;
    }

    // Request buffers
    struct v4l2_requestbuffers req;
    memset(&req, 0, sizeof(req));
    req.count = BUFFER_COUNT;
    req.type = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
    req.memory = V4L2_MEMORY_MMAP;

    if (ioctl(video_fd, VIDIOC_REQBUFS, &req) < 0) {
        perror("Failed to request buffers");
        close(video_fd);
        return -1;
    }

    // Map buffers
    struct v4l2_plane planes[1];
    for (int i = 0; i < BUFFER_COUNT; i++) {
        memset(&v4l2_buf, 0, sizeof(v4l2_buf));
        memset(planes, 0, sizeof(planes));
        v4l2_buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
        v4l2_buf.memory = V4L2_MEMORY_MMAP;
        v4l2_buf.index = i;
        v4l2_buf.m.planes = planes;
        v4l2_buf.length = 1;

        if (ioctl(video_fd, VIDIOC_QUERYBUF, &v4l2_buf) < 0) {
            perror("Failed to query buffer");
            close(video_fd);
            return -1;
        }

        buffers[i] = mmap(NULL, planes[0].length, PROT_READ | PROT_WRITE, MAP_SHARED, video_fd, planes[0].m.mem_offset);
        if (buffers[i] == MAP_FAILED) {
            perror("Failed to map buffer");
            close(video_fd);
            return -1;
        }

        // Queue buffer
        if (ioctl(video_fd, VIDIOC_QBUF, &v4l2_buf) < 0) {
            perror("Failed to queue buffer");
            close(video_fd);
            return -1;
        }
    }

    // Start streaming
    int type = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
    if (ioctl(video_fd, VIDIOC_STREAMON, &type) < 0) {
        perror("Failed to start streaming");
        close(video_fd);
        return -1;
    }

    return 0;
}

MppEncCfg mpp_enc_cfg(int fps, int bitrate) {
    MppEncCfg cfg;
    mpp_enc_cfg_init(&cfg);
    mpp_enc_cfg_set_s32(cfg, "prep:width", WIDTH);
    mpp_enc_cfg_set_s32(cfg, "prep:height", HEIGHT);
    mpp_enc_cfg_set_s32(cfg, "prep:format", MPP_FMT_BGR888);
    mpp_enc_cfg_set_s32(cfg, "prep:hor_stride", WIDTH * 3);
    mpp_enc_cfg_set_s32(cfg, "prep:ver_stride", HEIGHT);

    mpp_enc_cfg_set_s32(cfg, "rc:mode", MPP_ENC_RC_MODE_CBR);

    mpp_enc_cfg_set_s32(cfg, "rc:fps_in_flex", 0);            // Fixed input FPS
    mpp_enc_cfg_set_s32(cfg, "rc:fps_in_num", fps);            // Input FPS numerator
    mpp_enc_cfg_set_s32(cfg, "rc:fps_in_denorm", 1);          // Input FPS denominator
    mpp_enc_cfg_set_s32(cfg, "rc:fps_out_flex", 0);           // Fixed output FPS
    mpp_enc_cfg_set_s32(cfg, "rc:fps_out_num", fps);           // Output FPS numerator
    mpp_enc_cfg_set_s32(cfg, "rc:fps_out_denorm", 1);         // Output FPS denominator

    mpp_enc_cfg_set_s32(cfg, "rc:gop", 5);

    mpp_enc_cfg_set_u32(cfg, "rc:drop_mode", MPP_ENC_RC_DROP_FRM_DISABLED);
    mpp_enc_cfg_set_u32(cfg, "rc:drop_thd", 20);        /* 20% of max bps */
    mpp_enc_cfg_set_u32(cfg, "rc:drop_gap", 1);         /* Do not continuous drop frame */

    mpp_enc_cfg_set_s32(cfg, "rc:bps_target", bitrate);
    mpp_enc_cfg_set_s32(cfg, "rc:bps_max", bitrate * 17 / 16);
    mpp_enc_cfg_set_s32(cfg, "rc:bps_min", bitrate * 15 / 16);

    return cfg;
}

int mpp_init_wrapper(int fps, int bitrate) {
    MPP_RET ret = mpp_create(&mpp_ctx, &mpp_api);
    if (ret != MPP_OK) 
        return -1;

    ret = mpp_init(mpp_ctx, MPP_CTX_ENC, MPP_VIDEO_CodingAVC);
    if (ret != MPP_OK)
        return -1;
    
    MppPollType poll_type = MPP_POLL_BLOCK;
	ret = mpp_api->control(mpp_ctx, MPP_SET_OUTPUT_TIMEOUT,(MppParam)&poll_type);
    if (ret != MPP_OK) {
        mpp_destroy(mpp_ctx);
        return -1;
    }

    cfg = mpp_enc_cfg(fps, bitrate);
    ret = mpp_api->control(mpp_ctx, MPP_ENC_SET_CFG, cfg);
    if (ret != MPP_OK) {
        mpp_destroy(mpp_ctx);
        return -1;
    }

    // Get SPS and PPS information
    MppPacket extra_pkt = NULL;
    ret = mpp_api->control(mpp_ctx, MPP_ENC_GET_EXTRA_INFO, &extra_pkt);
    if (ret == MPP_OK && extra_pkt) {
        void *extra_data = mpp_packet_get_pos(extra_pkt);
        size_t extra_size = mpp_packet_get_length(extra_pkt);

        memcpy(sps_buffer, extra_data, extra_size);
        sps_size = extra_size;

        mpp_packet_deinit(&extra_pkt);
    } else {
        fprintf(stderr, "Failed to get SPS/PPS data from encoder\n");
    }

    return 0;
}

int get_sps(unsigned char *output, int *output_len) {
    if (sps_size == 0) {
        fprintf(stderr, "SPS data not available\n");
        return -1;
    }

    memcpy(output, sps_buffer, sps_size);
    *output_len = sps_size;

    return 0;
}

int capture_and_encode(unsigned char *output, int *output_len) {
    // Dequeue buffer
    struct v4l2_plane planes[1];
    memset(&v4l2_buf, 0, sizeof(v4l2_buf));
    memset(planes, 0, sizeof(planes));
    v4l2_buf.type = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
    v4l2_buf.memory = V4L2_MEMORY_MMAP;
    v4l2_buf.m.planes = planes;
    v4l2_buf.length = NUM_PLANES;

    if (ioctl(video_fd, VIDIOC_DQBUF, &v4l2_buf) < 0) {
        perror("Failed to dequeue buffer");
        return -1;
    }

    // Encode frame
    MppFrame frame;
    mpp_frame_init(&frame);
    mpp_frame_set_width(frame, WIDTH);
    mpp_frame_set_height(frame, HEIGHT);
    mpp_frame_set_fmt(frame, MPP_FMT_BGR888);

    MppBuffer frm_buf;
    mpp_buffer_get(NULL, &frm_buf, planes[0].length);
    memcpy(mpp_buffer_get_ptr(frm_buf), buffers[v4l2_buf.index], planes->bytesused);
    mpp_frame_set_buffer(frame, frm_buf);

    MPP_RET ret = mpp_api->encode_put_frame(mpp_ctx, frame);
    if (ret == MPP_OK) {
        MppPacket packet;
        ret = mpp_api->encode_get_packet(mpp_ctx, &packet);
        if (ret == MPP_OK) {
            *output_len = mpp_packet_get_length(packet);
            memcpy(output, mpp_packet_get_data(packet), *output_len);
            mpp_packet_deinit(&packet);
        }
    }

    mpp_frame_deinit(&frame);
    mpp_buffer_put(frm_buf);

    // Requeue buffer
    ioctl(video_fd, VIDIOC_QBUF, &v4l2_buf);
    return ret == MPP_OK ? 0 : -1;
}

void cleanup() {
    //if (cfg) mpp_enc_cfg_deinit(cfg);
    if (mpp_ctx) mpp_destroy(mpp_ctx);
    if (video_fd >= 0) close(video_fd);
}