#include "k230_vdec.h"
#include <iostream>
#include <unistd.h>
#include "mpi_vdec_api.h"
#include "mpi_sys_api.h"
#include "mpi_vb_api.h"

// VideoEncoder private method implementation
static k_payload_type _decoder_type_to_mpitype(DecoderType type) {
    switch (type) {
        case DecoderType::VDEC_H264: return K_PT_H264;
        case DecoderType::VDEC_H265: return K_PT_H265;
        case DecoderType::VDEC_JPEG: return K_PT_JPEG;
        default: return K_PT_H264;
    }
}

DecoderConfig::DecoderConfig(DecoderType dec_type, uint32_t in_buf_num , uint32_t out_buf_num, uint32_t w, uint32_t h)
    : type(dec_type), max_width(w), max_height(h), input_buf_num(in_buf_num),output_buf_num(out_buf_num) {}

int VideoDecoder::_setup_buffers() {
    // Create output buffer pool for vdec channel
    k_vb_pool_config pool_config;
    memset(&pool_config, 0, sizeof(pool_config));
    pool_config.blk_cnt = config_.output_buf_num;
    pool_config.blk_size = config_.max_width * config_.max_height * 2;// Frame
    pool_config.mode = VB_REMAP_MODE_NOCACHE;

    private_output_pool_id_ = kd_mpi_vb_create_pool(&pool_config);
    if (private_output_pool_id_ < 0) {
        std::cerr << "Failed to create vdec output buffer pool for channel " << channel_id_ << std::endl;
        return -1;
    }

    // Create input buffer pool for vdec channel
    memset(&pool_config, 0, sizeof(pool_config));
    pool_config.blk_cnt = config_.input_buf_num;
    pool_config.blk_size = config_.max_width * config_.max_height; // Stream
    pool_config.mode = VB_REMAP_MODE_NOCACHE;
    private_input_pool_id_ = kd_mpi_vb_create_pool(&pool_config);
    if (private_input_pool_id_ < 0) {
        std::cerr << "Failed to create vdec input buffer pool for channel " << channel_id_ << std::endl;
        kd_mpi_vb_destory_pool(private_output_pool_id_);
        private_output_pool_id_ = -1;
        return -1;
    }

    return 0;
}

void VideoDecoder::_release_buffers() {
    if (private_output_pool_id_ >= 0) {
        kd_mpi_vb_destory_pool(private_output_pool_id_);
        private_output_pool_id_ = -1;
    }

    if (private_input_pool_id_ >= 0) {
        kd_mpi_vb_destory_pool(private_input_pool_id_);
        private_input_pool_id_ = -1;
    }
}

// VideoDecoder public methods implementation
VideoDecoder::VideoDecoder(const DecoderConfig& config)
    : channel_id_(-1), state_(DESTROYED),config_(config),
    private_output_pool_id_(VB_INVALID_POOLID),private_input_pool_id_(VB_INVALID_POOLID) {
    memset(&input_stream_, 0, sizeof(input_stream_));
}

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

int VideoDecoder::create() {
    if (state_ != DESTROYED) {
        std::cerr << "Decoder already created" << std::endl;
        return -1;
    }

    // Get free channel
    kd_mpi_vdec_request_chn(&channel_id_);
    if (channel_id_ < 0) {
        std::cerr << "No available decoder channels" << std::endl;
        return -1;
    }
    std::cout << "Allocated decoder channel ID: " << channel_id_ << std::endl;

    // Setup buffer pools
    if (_setup_buffers() != 0) {
        channel_id_ = -1;
        return -1;
    }

    //attach pool vb
    kd_mpi_vdec_attach_vb_pool(channel_id_,private_output_pool_id_);

    // Configure decoder attributes
    k_vdec_chn_attr vdec_attr;
    memset(&vdec_attr, 0, sizeof(vdec_attr));
    vdec_attr.pic_width = config_.max_width;
    vdec_attr.pic_height = config_.max_height;
    vdec_attr.stream_buf_size = config_.max_width*config_.max_height;
    vdec_attr.type = _decoder_type_to_mpitype(config_.type);

    // Create decoder channel
    int ret = kd_mpi_vdec_create_chn(channel_id_, &vdec_attr);
    if (ret != 0) {
        std::cerr << "Failed to create decoder channel " << channel_id_ << std::endl;
        _release_buffers();
        channel_id_ = -1;
        return -1;
    }

    state_ = CREATED;
    return 0;
}

int VideoDecoder::start() {
    if (state_ != CREATED && state_ != STOPPED) {
        std::cerr << "Decoder not in correct state for starting" << std::endl;
        return -1;
    }

    int ret = kd_mpi_vdec_start_chn(channel_id_);
    if (ret != 0) {
        std::cerr << "Failed to start decoder channel " << channel_id_ << std::endl;
        return -1;
    }

    state_ = STARTED;
    return 0;
}

int VideoDecoder::send_stream(const uint8_t* data, uint32_t size, uint64_t timestamp, int timeout) {
    k_vb_blk_handle handle = 0;
    k_u32 stream_size = 0;

    if (state_ != STARTED) {
        std::cerr << "Decoder not started" << std::endl;
        return -1;
    }

    if (!data || size == 0) {
        std::cerr << "Invalid stream data" << std::endl;
        return -1;
    }

    stream_size = config_.max_width * config_.max_height;

    // Get vb from input pool
    while(true){
        handle = kd_mpi_vb_get_block(private_input_pool_id_, stream_size, NULL);
        if (handle == VB_INVALID_HANDLE) {
            //std::cerr << "Failed to get input buffer handle" << std::endl;
            usleep(1000*10);
        }
        else{
            break;
        }
    }

    // Map physical address to virtual address
    uint64_t phys_addr = kd_mpi_vb_handle_to_phyaddr(handle);
    void* vir_addr = kd_mpi_sys_mmap(phys_addr, stream_size);
    if (!vir_addr) {
        std::cerr << "Failed to map buffer address" << std::endl;
        kd_mpi_vb_release_block(handle);
        return -1;
    }

    // Copy stream data to buffer
    memcpy(vir_addr, data, size);

    // Unmap buffer
    kd_mpi_sys_munmap(vir_addr, stream_size);

    // Prepare stream structure
    input_stream_.end_of_stream = K_FALSE;
    input_stream_.len = size;
    input_stream_.phy_addr = phys_addr;
    input_stream_.pts = timestamp;

    // Send stream to decoder
    int ret = kd_mpi_vdec_send_stream(channel_id_, &input_stream_, timeout);
    if (ret != 0) {
        std::cerr << "Failed to send stream to decoder" << std::endl;
        kd_mpi_vb_release_block(handle);
        return -1;
    }

    // Release buffer handle
    kd_mpi_vb_release_block(handle);
    return 0;
}

int VideoDecoder::get_frame(k_video_frame_info& frame, int timeout) {
    k_vdec_supplement_info supplement;
    if (state_ != STARTED) {
        std::cerr << "Decoder not started" << std::endl;
        return -1;
    }

    return kd_mpi_vdec_get_frame(channel_id_, &frame, &supplement, timeout);
}

int VideoDecoder::release_frame(k_video_frame_info& frame) {
    if (state_ != STARTED) {
        std::cerr << "Decoder not started" << std::endl;
        return -1;
    }

    return kd_mpi_vdec_release_frame(channel_id_, &frame);
}

int VideoDecoder::stop() {
    if (state_ != STARTED) {
        std::cerr << "Decoder not running" << std::endl;
        return -1;
    }

    // Send end of stream signal
    k_vdec_stream eos_stream;
    memset(&eos_stream, 0, sizeof(eos_stream));
    eos_stream.end_of_stream = K_TRUE;
    eos_stream.len = 0;
    eos_stream.pts = 0;

    // Get a buffer for EOS signal
    // k_vb_blk_handle handle = kd_mpi_vb_get_block(private_input_pool_id_, config_.max_width * config_.max_height, "");
    // if (handle != VB_INVALID_HANDLE) {
    //     eos_stream.phy_addr = kd_mpi_vb_handle_to_phyaddr(handle);
    //     kd_mpi_vdec_send_stream(channel_id_, &eos_stream, -1);
    //     kd_mpi_vb_release_block(handle);
    // }

    k_vb_blk_handle handle;
    // 循环直到成功获取VB块并完成处理
    while (1) {
        handle = kd_mpi_vb_get_block(private_input_pool_id_,
                                    config_.max_width * config_.max_height,
                                    "");
        if (handle != VB_INVALID_HANDLE) {
            eos_stream.phy_addr = kd_mpi_vb_handle_to_phyaddr(handle);
            // 发送流数据
            kd_mpi_vdec_send_stream(channel_id_, &eos_stream, -1);
            // 释放VB块
            kd_mpi_vb_release_block(handle);
            // 成功处理后退出循环
            break;
        }
        usleep(1000*10); // 延迟1毫秒，根据实际情况调整
    }



    // Wait for all frames to be processed
    k_vdec_chn_status status;
    while (true) {
        kd_mpi_vdec_query_status(channel_id_, &status);
        if (status.end_of_stream) {
            break;
        }
        usleep(1000*10); // Sleep 1ms before rechecking
    }

    // Stop the decoder channel
    int ret = kd_mpi_vdec_stop_chn(channel_id_);
    if (ret != 0) {
        std::cerr << "Failed to stop decoder channel" << std::endl;
        return -1;
    }

    state_ = STOPPED;
    return 0;
}

int VideoDecoder::destroy() {
    if (state_ == DESTROYED) {
        return -1;
    }

    // Stop if still running
    if (state_ == STARTED) {
        stop();
    }

    // Destroy decoder channel
    if (channel_id_ >= 0) {
        kd_mpi_vdec_destroy_chn(channel_id_);
        kd_mpi_vdec_release_chn(channel_id_);
        // Release buffer pools
        _release_buffers();

        channel_id_ = -1;
    }

    state_ = DESTROYED;
    return 0;
}
