#include <opencv2/opencv.hpp>
#include <iostream>
#include <string>
#include <stdexcept>
#include <cstring>
#include <cmath>
#include <unistd.h>
#include "processor.hpp"

    VideoMemoryRaw::VideoMemoryRaw()  : data(nullptr), width(0), height(0), channels(3), frameCount(0), fps(25.0) {
    }

    int Processor::idx(int y, int x, int width) {
        return y * width + x;
    }

    float* Processor::allocateFloatBuffer(int size) {
        float* p = new float[size];
	for (int i=0;i<size;i++)
		p[i]=0;
        return p;
    }

    uint8_t* Processor::allocateU8Buffer(int size) {
        uint8_t* p = new uint8_t[size];
	for (int i=0;i<size;i++)
		p[i]=0;
        return p;
    }

    void Processor::extractChannel(uint8_t* bgr,
                        int width,
                        int height,
                        int channel,
                        float* out) {
        const int pixels = width * height;
        for (int i = 0; i < pixels; ++i) {
            out[i] = (float) bgr[3 * i + channel];
        }
    }

    void Processor::interleaveChannels(uint8_t* B,
                            uint8_t* G,
                            uint8_t* R,
                            int width,
                            int height,
                            uint8_t* outBGR) {
        int pixels = width * height;
        for (int i = 0; i < pixels; ++i) {
            outBGR[3 * i + 0] = B[i];
            outBGR[3 * i + 1] = G[i];
            outBGR[3 * i + 2] = R[i];
        }
    }

//public:

    VideoMemoryRaw*  Processor::loadVideoToMemoryRowMajor(const std::string& inputPath) {
    cv::VideoCapture cap(inputPath);
    if (!cap.isOpened()) {
        std::cout << "Could not open input video: " << inputPath << std::endl;
        _exit(1);
    }

    VideoMemoryRaw* vm = new VideoMemoryRaw();

    vm->width = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_WIDTH));
    vm->height = static_cast<int>(cap.get(cv::CAP_PROP_FRAME_HEIGHT));
    vm->channels = 3;
    vm->fps = cap.get(cv::CAP_PROP_FPS);
    if (vm->fps <= 0.0) {
        vm->fps = 25.0;
    }

    const int frameBytes = vm->width * vm->height * vm->channels;
    int countedFrames = 0;
    cv::Mat frame;

    while (cap.read(frame)) {
        if (frame.empty()) {
            break;
        }
        countedFrames++;
    }

    if (countedFrames == 0) {
        cap.release();
        delete vm;
        std::cout << "No frames found in video." << std::endl;
        _exit(1);
    }

    vm->frameCount = countedFrames;
    vm->data = new uint8_t[(size_t)vm->frameCount * frameBytes];

    cap.release();
    cap.open(inputPath);

    if (!cap.isOpened()) {
        delete[] vm->data;
        delete vm;
        std::cout << "Could not reopen input video: " << inputPath << std::endl;
        _exit(1);
    }

    int currentIndex = 0;
    while (cap.read(frame)) {
        if (frame.empty()) {
            break;
        }

        if (frame.channels() != 3 || frame.depth() != CV_8U) {
            cap.release();
            delete[] vm->data;
            delete vm;
            std::cout << "Each frame must be an 8-bit BGR image." << std::endl;
            _exit(1);
        }

        //if (!frame.isContinuous()) {
        //    frame = frame.clone();
        //}

        std::memcpy(vm->data + (size_t)currentIndex * frameBytes, frame.data, frameBytes);
        currentIndex++;
    }

    cap.release();
    return vm;
}

void Processor::freeVideoMemory(VideoMemoryRaw* vm) {
    if (vm != nullptr) {
        if (vm->data != nullptr) {
            delete[] vm->data;
            vm->data = nullptr;
        }
        delete vm;
    }
}

uint8_t* Processor::getFramePointer(VideoMemoryRaw* vm, int frameIndex) {
    if (vm == nullptr || vm->data == nullptr) {
        std::cout << "Video memory is empty." << std::endl;
        _exit(1);
    }

    if (frameIndex < 0 || frameIndex >= vm->frameCount) {
        std::cout << "Frame index out of range: " << frameIndex << std::endl;
        _exit(1);
    }

    const int frameBytes = vm->width * vm->height * vm->channels;
    return vm->data + (size_t)frameIndex * frameBytes;
}

void Processor::saveVideoFromMemory(VideoMemoryRaw* vm,
                         const std::string& outputPath = "output.mp4") {
    if (vm == nullptr || vm->data == nullptr || vm->frameCount <= 0) {
        std::cout << "No video frames in memory." << std::endl;
        _exit(1);
    }

    if (vm->width <= 0 || vm->height <= 0 || vm->channels != 3) {
        std::cout << "Invalid video memory metadata." << std::endl;
        _exit(1);
    }

    double fps = vm->fps;
    if (fps <= 0.0) {
        fps = 25.0;
    }

    cv::VideoWriter writer;
    int fourcc = cv::VideoWriter::fourcc('a', 'v', 'c', '1');

    writer.open(outputPath, fourcc, fps, cv::Size(vm->width, vm->height), true);

    if (!writer.isOpened()) {
        std::cout << "Could not open output video: " << outputPath << std::endl;
        _exit(1);
    }

    const int frameBytes = vm->width * vm->height * vm->channels;

    for (int i = 0; i < vm->frameCount; ++i) {
        uint8_t* framePtr = vm->data + (size_t)i * frameBytes;

        cv::Mat frame(vm->height, vm->width, CV_8UC3, framePtr);

        writer.write(frame);
    }

    writer.release();

    std::cout << "Saved video from memory to: " << outputPath << std::endl;
}


void  Processor::saveVideoFromMemoryRange(VideoMemoryRaw* vm,
                              int startFrame,
                              int endFrame,
                              const std::string& outputPath = "output.mp4") {
    if (vm == nullptr || vm->data == nullptr || vm->frameCount <= 0) {
        std::cout << "No video frames in memory." << std::endl;
        _exit(1);
    }
    if (vm->width <= 0 || vm->height <= 0 || vm->channels != 3) {
        std::cout << "Invalid video memory metadata." << std::endl;
        _exit(1);
    }

    if (startFrame < 0) startFrame = 0;
    if (endFrame >= vm->frameCount) endFrame = vm->frameCount - 1;

    if (startFrame > endFrame) {
        std::cout << "Invalid frame range." << std::endl;
        _exit(1);
    }

    double fps = vm->fps;
    if (fps <= 0.0) {
        fps = 25.0;
    }

    cv::VideoWriter writer;
    int fourcc = cv::VideoWriter::fourcc('a', 'v', 'c', '1');

    writer.open(outputPath, fourcc, fps, cv::Size(vm->width, vm->height), true);

    if (!writer.isOpened()) {
        std::cout << "Could not open output video: " << outputPath << std::endl;
        _exit(1);
    }

    const int pixels = vm->width * vm->height;
    const int frameBytes = pixels * vm->channels;

    for (int i = startFrame; i <= endFrame; ++i) {
        uint8_t* framePtr = vm->data + (size_t)i * frameBytes;
        cv::Mat frame(vm->height, vm->width, CV_8UC3, framePtr);
        writer.write(frame);
    }

    writer.release();

    std::cout << "Saved video frame range [" << startFrame
              << ", " << endFrame << "] to: " << outputPath << std::endl;
}

void Processor::saveVideoFromMemory(VideoMemoryRaw* vm,
                                                    const std::string& outputPath,
                                                    int selectedChannel) {
    if (vm == nullptr || vm->data == nullptr || vm->frameCount <= 0) {
        std::cout << "No video frames in memory." << std::endl;
        _exit(1);
    }

    if (vm->width <= 0 || vm->height <= 0 || vm->channels != 3) {
        std::cout << "Invalid video memory metadata." << std::endl;
        _exit(1);
    }

    if (selectedChannel < 0 || selectedChannel > 2) {
        std::cout << "selectedChannel must be 0 (B), 1 (G), or 2 (R)." << std::endl;
        _exit(1);
    }
    if (outputPath.empty()) {
	std::cout << " Empty filename to store data !" << std::endl;
	_exit(1);
    }

    std::string finalOutput = outputPath;

    double fps = vm->fps;
    if (fps <= 0.0) {
        fps = 25.0;
    }

    cv::VideoWriter writer;
    int fourcc = cv::VideoWriter::fourcc('a', 'v', 'c', '1');
    writer.open(finalOutput, fourcc, fps, cv::Size(vm->width, vm->height), true);

    if (!writer.isOpened()) {
        std::cout << "Could not open output video: " << finalOutput << std::endl;
        _exit(1);
    }

    const int pixels = vm->width * vm->height;
    const int frameBytes = pixels * vm->channels;

    uint8_t* outFrameData = new uint8_t[frameBytes];

    for (int f = 0; f < vm->frameCount; ++f) {
        uint8_t* srcFrame = vm->data + static_cast<size_t>(f) * frameBytes;

        for (int i = 0; i < pixels; ++i) {
            //outFrameData[3 * i + 0] = 255;
            //outFrameData[3 * i + 1] = 255;
            //outFrameData[3 * i + 2] = 255;
            outFrameData[3 * i + 0] = 0;
            outFrameData[3 * i + 1] = 0;
            outFrameData[3 * i + 2] = 0;

            outFrameData[3 * i + selectedChannel] = srcFrame[3 * i + selectedChannel];
        }

        cv::Mat outFrame(vm->height, vm->width, CV_8UC3, outFrameData);
        writer.write(outFrame);
    }

    delete[] outFrameData;
    writer.release();

    std::cout << "Saved selective-channel video to: " << finalOutput << std::endl;
}

void Processor::saveThreeChannelVideosFromMemory(
    VideoMemoryRaw* vm,
    const std::string& blueOutputPath=0,
    const std::string& greenOutputPath=0,
    const std::string& redOutputPath=0) {

    if (vm == nullptr || vm->data == nullptr || vm->frameCount <= 0) {
        std::cout << "No video frames in memory." << std::endl;
        _exit(1);
    }

    if (vm->width <= 0 || vm->height <= 0 || vm->channels != 3) {
        std::cout << "Invalid video memory metadata." << std::endl;
        _exit(1);
    }

    double fps = vm->fps;
    if (fps <= 0.0) {
        fps = 25.0;
    }

    const std::string bluePath  = blueOutputPath.empty()  ? "blue_response.mp4"  : blueOutputPath;
    const std::string greenPath = greenOutputPath.empty() ? "green_response.mp4" : greenOutputPath;
    const std::string redPath   = redOutputPath.empty()   ? "red_response.mp4"   : redOutputPath;

    int fourcc = cv::VideoWriter::fourcc('a', 'v', 'c', '1');

    cv::VideoWriter blueWriter;
    cv::VideoWriter greenWriter;
    cv::VideoWriter redWriter;

    blueWriter.open(bluePath, fourcc, fps, cv::Size(vm->width, vm->height), true);
    greenWriter.open(greenPath, fourcc, fps, cv::Size(vm->width, vm->height), true);
    redWriter.open(redPath, fourcc, fps, cv::Size(vm->width, vm->height), true);

    if (!blueWriter.isOpened()) {
        std::cout << "Could not open output video: " << bluePath << std::endl;
        _exit(1);
    }
    if (!greenWriter.isOpened()) {
        std::cout << "Could not open output video: " << greenPath << std::endl;
        _exit(1);
    }
    if (!redWriter.isOpened()) {
        std::cout << "Could not open output video: " << redPath << std::endl;
        _exit(1);
    }

    const int pixels = vm->width * vm->height;
    const int frameBytes = pixels * vm->channels;

    uint8_t* blueFrameData = new uint8_t[frameBytes];
    uint8_t* greenFrameData = new uint8_t[frameBytes];
    uint8_t* redFrameData = new uint8_t[frameBytes];

    for (int f = 0; f < vm->frameCount; ++f) {
        uint8_t* srcFrame = vm->data + (size_t)f * frameBytes;

        for (int i = 0; i < pixels; ++i) {
            const uint8_t b = srcFrame[3 * i + 0];
            const uint8_t g = srcFrame[3 * i + 1];
            const uint8_t r = srcFrame[3 * i + 2];
	    //Yellow
            //blueFrameData[3 * i + 0] = b;
            //blueFrameData[3 * i + 1] = 255;
            //blueFrameData[3 * i + 2] = 255;
            blueFrameData[3 * i + 0] = 255;
            blueFrameData[3 * i + 1] = 255;
            blueFrameData[3 * i + 2] = b;
	    //Pink
            //greenFrameData[3 * i + 0] = 255;
            //greenFrameData[3 * i + 1] = g;
            //greenFrameData[3 * i + 2] = 255;
            greenFrameData[3 * i + 0] = 255;
            greenFrameData[3 * i + 1] = g;
            greenFrameData[3 * i + 2] = 255;
	    //Cyan
            //redFrameData[3 * i + 0] = 255;
            //redFrameData[3 * i + 1] = 255;
            //redFrameData[3 * i + 2] = r;
            redFrameData[3 * i + 0] = r;
            redFrameData[3 * i + 1] = 255;
            redFrameData[3 * i + 2] = 255;
        }

        cv::Mat blueFrame(vm->height, vm->width, CV_8UC3, blueFrameData);
        cv::Mat greenFrame(vm->height, vm->width, CV_8UC3, greenFrameData);
        cv::Mat redFrame(vm->height, vm->width, CV_8UC3, redFrameData);

        blueWriter.write(blueFrame);
        greenWriter.write(greenFrame);
        redWriter.write(redFrame);
    }

    delete[] blueFrameData;
    delete[] greenFrameData;
    delete[] redFrameData;

    blueWriter.release();
    greenWriter.release();
    redWriter.release();

    std::cout << "Saved blue response video to: " << bluePath << std::endl;
    std::cout << "Saved green response video to: " << greenPath << std::endl;
    std::cout << "Saved red response video to: " << redPath << std::endl;
}
//};