#include <fstream>
#include <iostream>
#include <vector>
#include <string>
#include <unistd.h>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <atomic>
#include <algorithm>
#include <cstdlib>  // For mkstemp
#include <iomanip>
#include <queue>

#include <k230_vdec.h>
#include <k230_venc.h>

// Size of the read buffer for each read operation
const size_t READ_BUFFER_SIZE = 40960;

std::atomic<bool> g_decoder_running(true);
std::atomic<bool> g_encoder_running(true);
std::atomic<bool> g_writer_running(true);

// Helper function to get file extension (lowercase)
static std::string get_file_extension(const std::string& file_path) {
    size_t dot_pos = file_path.find_last_of('.');
    if (dot_pos == std::string::npos || dot_pos == file_path.length() - 1) {
        return ""; // No extension
    }
    std::string ext = file_path.substr(dot_pos + 1);
    std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
    return ext;
}

// Helper function to determine decoder type from file extension
static DecoderType get_decoder_type_from_extension(const std::string& ext) {
    if (ext == "h264" || ext == "264") {
        return VDEC_H264;
    } else if (ext == "h265" || ext == "265" || ext == "hevc") {
        return VDEC_H265;
    } else if (ext == "mjpeg" || ext == "mjpg") {
        return VDEC_JPEG;
    }
    return VDEC_H264;
}

// Frame processing thread function - handles decoding and sending to encoder
static void decoded_frame_encoder_feeder_thread(VideoDecoder* decoder, VideoEncoder* encoder)
{
    if (!decoder || !encoder) {
        std::cerr << "Error: Invalid decoder or encoder pointer in processing thread" << std::endl;
        return;
    }

    int frame_count = 0;

    while (g_decoder_running && g_encoder_running) {
        k_video_frame_info frame;
        int ret = decoder->get_frame(frame, 100);
        if (ret == 0) {
            frame_count++;
            if (frame_count % 100 == 0){
                printf("[%d]:decoded frame,width:%d,height:%d\n",
                    frame_count,
                    frame.v_frame.width,
                    frame.v_frame.height);
            }

            // Send frame to encoder with mutex protection
            if (encoder->send_frame(&frame, -1) != 0) {
                std::cerr << "Warning: Failed to send frame to encoder" << std::endl;
            }

            // Release original decoded frame buffer
            decoder->release_frame(frame);
        }
    }

    std::cout << "Decoded frame encoder feeder thread exited. Total decoded frames: " << frame_count << std::endl;
}

// Stream writing thread - handles getting encoded stream and writing to file
static void encoded_stream_file_writer_thread(VideoEncoder* encoder, const std::string& output_file)
{
    if (!encoder) {
        std::cerr << "Error: Invalid encoder pointer in writing thread" << std::endl;
        return;
    }

    // Open output file for writing encoded data
    std::ofstream out_file(output_file, std::ios::binary);
    if (!out_file.is_open()) {
        std::cerr << "Error: Unable to open output file: " << output_file << std::endl;
        g_writer_running = false;
        return;
    }

    int encode_count = 0;

    while (g_writer_running) {
        // Check if we should exit
        if (!g_decoder_running && !g_encoder_running) {
            break;
        }

        // Get encoded stream with mutex protection
        EncodedStream stream;
        int enc_ret = encoder->get_stream(stream, 100);  // Short timeout to prevent blocking
        if (enc_ret == 0) {
            for (uint32_t i = 0;i < stream.packet_count;i ++){
                if (stream.packets[i].type != K_VENC_HEADER)
                    encode_count++;

                // Write encoded data to file
                out_file.write(reinterpret_cast<const char*>(stream.packets[i].data), stream.packets[i].size);
                if (encode_count % 100 == 0 ){
                    printf("[%d]:encoded frame,size:%d,type:%d\n",encode_count,stream.packets[i].size,stream.packets[i].type);
                }
            }

            // Release the stream buffer back to encoder
            encoder->release_stream(stream);
        }
    }

    // After main processing, check for any remaining encoded streams
    while (true) {
        EncodedStream stream;
        int enc_ret = encoder->get_stream(stream, 0);  // Non-blocking check
        if (enc_ret != 0) {
            break;
        }

        for (uint32_t i = 0;i < stream.packet_count;i ++){
            if (stream.packets[i].type != K_VENC_HEADER)
                encode_count++;

            // Write encoded data to file
            out_file.write(reinterpret_cast<const char*>(stream.packets[i].data), stream.packets[i].size);
            if (encode_count % 100 == 0 ){
                printf("[%d]:encoded frame,size:%d,type:%d\n",encode_count,stream.packets[i].size,stream.packets[i].type);
            }
        }

        encoder->release_stream(stream);
    }

    // Cleanup
    out_file.close();

    std::cout << "Encoded stream file writer thread exited. Total encoded streams: " << encode_count << std::endl;
}

// Modified function with separated encoding and writing threads
static int file_video_transcoder(const char* input_file_path, const char* output_path) {
    int ret = -1;

    if (!input_file_path || !output_path) {
        std::cerr << "Error: File path or output path is empty" << std::endl;
        return -1;
    }

    // Get file extension and determine decoder type
    std::string ext = get_file_extension(input_file_path);
    DecoderType dec_type = get_decoder_type_from_extension(ext);

    if (dec_type == VDEC_UNKNOWN) {
        std::cerr << "Error: Unsupported file format (extension: ." << ext << ")" << std::endl;
        return -1;
    }
    std::cout << "Detected file format: ." << ext << ", using decoder type: " << dec_type << std::endl;

    try {
        // Configure decoder parameters
        DecoderConfig config(dec_type);

        // Create decoder instance
        VideoDecoder decoder(config);

        // Initialize decoder
        if (decoder.create() != 0) {
            std::cerr << "Error: Failed to create decoder" << std::endl;
            return -1;
        }

        // Start decoder
        if (decoder.start() != 0) {
            std::cerr << "Error: Failed to start decoder" << std::endl;
            decoder.destroy();
            return -1;
        }

        // Configure encoder - using H264 as default
        EncoderConfig enc_config(VENC_H264, VENC_PROFILE_H264_MAIN, 1920, 1080, 4, 2000, 30, 30, 30);
        VideoEncoder encoder(enc_config);

        // Initialize encoder
        if (encoder.create() != 0) {
            std::cerr << "Error: Failed to create encoder" << std::endl;
            decoder.stop();
            decoder.destroy();
            return -1;
        }

        if (encoder.start() != 0) {
            std::cerr << "Error: Failed to start encoder" << std::endl;
            encoder.destroy();
            decoder.stop();
            decoder.destroy();
            return -1;
        }

        // Create and start threads
        std::thread processing_thread(decoded_frame_encoder_feeder_thread, &decoder, &encoder);
        std::thread writing_thread(encoded_stream_file_writer_thread, &encoder, output_path);

        // Open input file
        std::ifstream file(input_file_path, std::ios::binary);
        if (!file.is_open()) {
            std::cerr << "Error: Unable to open file: " << input_file_path << std::endl;
            // Notify threads to exit
            g_decoder_running = false;
            g_encoder_running = false;
            g_writer_running = false;
            processing_thread.join();
            writing_thread.join();
            decoder.stop();
            decoder.destroy();
            encoder.stop();
            encoder.destroy();
            return -1;
        }

        //get total file size
        file.seekg(0, std::ios::end);
        size_t file_total_size = file.tellg();
        file.seekg(0, std::ios::beg);

        // Read file and send data to decoder (main thread)
        std::vector<uint8_t> read_buffer(READ_BUFFER_SIZE);
        size_t total_read = 0;

        while (file && g_decoder_running) {
            file.read(reinterpret_cast<char*>(read_buffer.data()), READ_BUFFER_SIZE);
            size_t bytes_read = file.gcount();

            if (bytes_read == 0) {
                break;
            }

            total_read += bytes_read;
            std::cout << "Read data: " << bytes_read << " bytes, total: " << total_read
                << " bytes (" << std::fixed << std::setprecision(2)
                << (total_read * 100.0 / file_total_size) << "%)" << std::endl;

            // Send data to decoder
            ret = decoder.send_stream(read_buffer.data(), bytes_read);
            if (ret != 0) {
                std::cerr << "Error: Failed to send stream data to decoder" << std::endl;
                break;
            }
        }

        std::cout << "File reading completed, total size: " << total_read << " bytes" << std::endl;
        sleep(3);

        // Wait for all frames to be processed
        g_decoder_running = false;

        // Wait for processing thread to complete
        if (processing_thread.joinable()) {
            processing_thread.join();
        }

        // Signal encoder completion and wait for writer thread
        g_encoder_running = false;
        g_writer_running = false;
        if (writing_thread.joinable()) {
            writing_thread.join();
        }

        // Stop and destroy decoder and encoder
        decoder.stop();
        decoder.destroy();

        encoder.stop();
        encoder.destroy();

        // Reset global variables
        g_decoder_running = true;
        g_encoder_running = true;
        g_writer_running = true;

        return 0; // Success

    } catch (const std::exception& e) {
        std::cerr << "Exception occurred: " << e.what() << std::endl;
        g_decoder_running = false;
        g_encoder_running = false;
        g_writer_running = false;
        return -1;
    } catch (...) {
        std::cerr << "Unknown exception occurred" << std::endl;
        g_decoder_running = false;
        g_encoder_running = false;
        g_writer_running = false;
        return -1;
    }
}

void setup()
{
    // Decode test.264 and encode to output.264
    file_video_transcoder("/data/test.264", "/data/output.264");
}

void loop()
{

}