#include "k230_venc.h"
#include <iostream>
#include "mpi_venc_api.h"
#include "mpi_sys_api.h"
#include "mpi_vb_api.h"

#define K_INVALID_CHANNEL             (-1U)

// EncoderConfig constructor
EncoderConfig::EncoderConfig(EncoderType enc_type, k_venc_profile prof, uint32_t w, uint32_t h, uint32_t buf_num,
              uint32_t br, uint32_t g,
              uint32_t sfps, uint32_t dfps,
              uint32_t jq)
    : type(enc_type), profile(prof), width(w), height(h),
      bitrate(br), gop(g), src_fps(sfps), dst_fps(dfps),
      jpeg_quality(jq), output_buffers(buf_num) {}

// VideoEncoder private method implementation
static k_payload_type _encoder_type_to_mpitype(EncoderType type) {
    switch (type) {
        case EncoderType::VENC_H264: return K_PT_H264;
        case EncoderType::VENC_H265: return K_PT_H265;
        case EncoderType::VENC_JPEG: return K_PT_JPEG;
        default: return K_PT_H264;
    }
}

int VideoEncoder::_setup_output_buffers() {
    if (config_.output_buffers == 0) return -1;

    k_vb_pool_config pool_config;
    memset(&pool_config, 0, sizeof(pool_config));
    pool_config.blk_cnt = config_.output_buffers;
    pool_config.blk_size = (config_.width * config_.height / 2 + 0xfff) & ~0xfff;
    pool_config.mode = VB_REMAP_MODE_NOCACHE;
    private_pool_id_ = kd_mpi_vb_create_pool(&pool_config);
    printf("%s poolid %d\n", __func__, private_pool_id_);

    return kd_mpi_venc_attach_vb_pool(channel_id_, private_pool_id_);
}

void VideoEncoder::_create_encoder_attributes(k_venc_chn_attr& venc_attr) {
    memset(&venc_attr, 0, sizeof(venc_attr));

    // Set encoder basic attributes
    venc_attr.venc_attr.type = _encoder_type_to_mpitype(config_.type);
    venc_attr.venc_attr.pic_width = config_.width;
    venc_attr.venc_attr.pic_height = config_.height;
    venc_attr.venc_attr.profile = config_.profile;

    // Set bitrate control
    if (config_.type == EncoderType::VENC_JPEG) {
        venc_attr.rc_attr.rc_mode = K_VENC_RC_MODE_MJPEG_FIXQP;
        venc_attr.rc_attr.mjpeg_fixqp.src_frame_rate = config_.src_fps;
        venc_attr.rc_attr.mjpeg_fixqp.dst_frame_rate = config_.dst_fps;
        venc_attr.rc_attr.mjpeg_fixqp.q_factor = config_.jpeg_quality;
    } else {
        venc_attr.rc_attr.rc_mode = K_VENC_RC_MODE_CBR;
        venc_attr.rc_attr.cbr.gop = config_.gop;
        venc_attr.rc_attr.cbr.stats_time = 0;
        venc_attr.rc_attr.cbr.src_frame_rate = config_.src_fps;
        venc_attr.rc_attr.cbr.dst_frame_rate = config_.dst_fps;
        venc_attr.rc_attr.cbr.bit_rate = config_.bitrate;
    }
}

VideoEncoder::VideoEncoder(const EncoderConfig& config)
    : channel_id_(K_INVALID_CHANNEL), state_(EncoderState::DESTROYED), private_pool_id_(VB_INVALID_POOLID),
      config_(config) {
    memset(&output_, 0, sizeof(output_));
}

VideoEncoder::~VideoEncoder() {
    destroy();
}

int VideoEncoder::create() {
    k_s32 ret;
    if (state_ != EncoderState::DESTROYED) {
        printf("Encoder already created\n");
        return -1;
    }

    kd_mpi_venc_request_chn(&channel_id_);
    if (channel_id_ < 0 || channel_id_ >= VENC_MAX_CHN_NUMS) {
        printf("Failed to request encoder channel\n");
        return -1;
    }

    // Set up output buffers
    _setup_output_buffers();

    // Create encoder attributes
    k_venc_chn_attr venc_attr;
    _create_encoder_attributes(venc_attr);

    // Create encoding channel
    ret = kd_mpi_venc_create_chn(channel_id_, &venc_attr);
    if (ret != 0) {
        printf("Failed to create encoder channel\n");
        return -1;
    }

    // Enable IDR for H264/H265
    if (config_.type == EncoderType::VENC_H264 || config_.type == EncoderType::VENC_H265) {
        ret = kd_mpi_venc_enable_idr(channel_id_, K_TRUE);
        if (ret != 0) {
            kd_mpi_venc_destroy_chn(channel_id_);
            printf("Failed to enable IDR\n");
            return -1;
        }
    }

    state_ = EncoderState::CREATED;

    return 0;
}

int VideoEncoder::start() {
    k_s32 ret;
    if (state_ != EncoderState::CREATED && state_ != EncoderState::STOPPED) {
        printf("Encoder not in correct state for starting\n");
        return -1;
    }

    ret = kd_mpi_venc_start_chn(channel_id_);
    if (ret != 0) {
        printf("Failed to start encoder\n");
        return -1;
    }

    state_ = EncoderState::STARTED;

    return 0;
}

int VideoEncoder::send_frame(k_video_frame_info* frame, int timeout) {
    if (state_ != EncoderState::STARTED) {
        printf("Encoder not started\n");
        return -1;
    }

    return kd_mpi_venc_send_frame(channel_id_, frame, timeout);
}

int VideoEncoder::get_stream(EncodedStream& stream, int timeout) {
    k_s32 ret;

    if (state_ != EncoderState::STARTED) {
        printf("Encoder not started\n");
        return -1;
    }

    k_venc_chn_status status;
    memset(&status, 0, sizeof(status));

    ret = kd_mpi_venc_query_status(channel_id_, &status);
    if (ret != 0) {
        printf("Failed to query encoder status\n");
        return -1;
    }

    uint32_t pack_count = (status.cur_packs > 0) ? status.cur_packs : 1;
    output_.pack_cnt = pack_count;
    output_.pack = venc_pack_;

    ret = kd_mpi_venc_get_stream(channel_id_, &output_, timeout);
    if (ret != 0) {
        return -1;
    }

    stream.packets.resize(pack_count);
    stream.packet_count = pack_count;
    for (uint32_t i = 0; i < pack_count; i++){
        stream.packets[i].data = kd_mpi_sys_mmap(output_.pack[i].phys_addr, output_.pack[i].len);
        stream.packets[i].phys_addr = output_.pack[i].phys_addr;
        stream.packets[i].size = output_.pack[i].len;
        stream.packets[i].type = output_.pack[i].type;
        stream.packets[i].pts = output_.pack[i].pts;
    }

    return 0;
}

int VideoEncoder::release_stream(EncodedStream& stream) {
    k_s32 ret;

    if (state_ != EncoderState::STARTED) {
        printf("Encoder not started\n");
        return -1;
    }

    for (uint32_t i = 0; i < stream.packet_count; i++) {
        if (stream.packets[i].data != nullptr) {
            ret = kd_mpi_sys_munmap(stream.packets[i].data, stream.packets[i].size);
            if (ret != 0) {
                printf("Failed to unmap stream data\n");
            }
            stream.packets[i].data = nullptr;
        }
    }

    ret = kd_mpi_venc_release_stream(channel_id_, &output_);
    if (ret != 0) {
        printf("Failed to release stream\n");
    }

    return 0;
}

int VideoEncoder::stop() {
    k_s32 ret;
    if (state_ != EncoderState::STARTED) {
        return -1;
    }

    ret = kd_mpi_venc_stop_chn(channel_id_);
    if (ret != 0) {
        throw std::runtime_error("Failed to stop encoder");
    }

    state_ = EncoderState::STOPPED;

    return 0;
}

int VideoEncoder::destroy() {
    k_s32 ret;
    if (state_ == EncoderState::DESTROYED) {
        return -1;
    }

    if (state_ == EncoderState::STARTED) {
        stop();
    }

    ret = kd_mpi_venc_detach_vb_pool(channel_id_);
    if (ret != 0) {
        printf("Failed to detach VB pool\n");
    }

    ret = kd_mpi_venc_destroy_chn(channel_id_);
    if (ret != 0) {
        printf("Failed to destroy encoder\n");
    }

    kd_mpi_venc_release_chn(channel_id_);
    channel_id_ = K_INVALID_CHANNEL;

    if (private_pool_id_ != VB_INVALID_POOLID) {
        kd_mpi_vb_destory_pool(private_pool_id_);
        private_pool_id_ = VB_INVALID_POOLID;
    }

    state_ = EncoderState::DESTROYED;

    return 0;
}