#include <stdexcept>
#include <vector>
#include <string>
#include <Processing.NDI.Lib.h>
#include "V210.hpp"

namespace dmf {
class NDIHelper {
    public:

        int xres = 0, yres = 0, frame_D = 0, frame_N = 0, stride = 0;

        NDIHelper() {
            if (!NDIlib_is_supported_CPU()) {
                throw std::runtime_error("CPU is not sufficient for NDI");
            }
            if (!NDIlib_initialize()) {
                throw std::runtime_error("NDI lib init failed");
            }
        }

        ~NDIHelper() {
            if (pNDI_recv) NDIlib_recv_destroy(pNDI_recv);
            if (pNDI_find) NDIlib_find_destroy(pNDI_find);
            NDIlib_destroy();
        }

        void find_sources(std::vector<std::string>* sources, u_int32_t timeout_ms) {
            pNDI_find = NDIlib_find_create_v2();
            if (!pNDI_find) {
                throw std::runtime_error("Cannot create NDI finder");
            }

            const int max_attempts = 10;
            uint32_t sources_amount = 0;
            const NDIlib_source_t* p_sources = nullptr;
            for (int attempt = 0; !sources_amount && attempt < max_attempts; ++attempt) {
                NDIlib_find_wait_for_sources(pNDI_find, timeout_ms);
                p_sources = NDIlib_find_get_current_sources(pNDI_find, &sources_amount);
            }
            if (!sources_amount) {
                NDIlib_find_destroy(pNDI_find);
                pNDI_find = nullptr;
                throw std::runtime_error("No NDI sources found after timeout");
            }

            // Copy while finder is alive: p_sources points into finder-owned memory
            for (uint32_t i = 0; i < sources_amount; ++i) {
                sources->push_back(p_sources[i].p_ndi_name);
                cached_names.emplace_back(p_sources[i].p_ndi_name);
                cached_urls.emplace_back(p_sources[i].p_url_address);
            }
            cached_sources.reserve(cached_names.size());
            for (uint32_t i = 0; i < cached_names.size(); ++i) {
                cached_sources.push_back({cached_names[i].c_str(), cached_urls[i].c_str()});
            }

            NDIlib_find_destroy(pNDI_find);
            pNDI_find = nullptr;
        }

        void select_source(uint32_t source_num) {
            if (cached_sources.empty()) {
                throw std::runtime_error("0 sources found");
            } else if (source_num >= cached_sources.size()) {
                throw std::runtime_error("Source_num bigger that sources amount");
            }

            pNDI_recv = NDIlib_recv_create_v3();
            if (!pNDI_recv) {
                throw std::runtime_error("Cannot create NDI recieve instance");
            }
            NDIlib_recv_connect(pNDI_recv, &cached_sources[source_num]);
        }

        void get_source_info(uint32_t source_num) {
            NDIlib_video_frame_v2_t video_frame;
            NDIlib_frame_type_e frame_type;
            bool is_got_info = false;
            while(!is_got_info)
            {
                frame_type = NDIlib_recv_capture_v3(pNDI_recv, &video_frame, nullptr, nullptr, 1000);
                switch(frame_type)
                {
                    case NDIlib_frame_type_video:
                        is_got_info = true;
                        xres = video_frame.xres;
                        yres = video_frame.yres;
                        frame_D = video_frame.frame_rate_D;
                        frame_N = video_frame.frame_rate_N;
                        fourCC = video_frame.FourCC;
                        stride = video_frame.line_stride_in_bytes;
                        if (stride == 0) {
                            stride = xres * get_bytes_per_pixel(fourCC);
                        }
                        break;
                    case NDIlib_frame_type_error:
                        is_got_info = true;
                        throw std::runtime_error("Selected NDI source is lost");
                        break;
                }
            }
            NDIlib_recv_free_video_v2(pNDI_recv, &video_frame);
        }

        int get_bytes_per_pixel(NDIlib_FourCC_video_type_e fourCC) {
            switch (fourCC) {
                case NDIlib_FourCC_video_type_UYVY: // Standard 8-bit YUV 4:2:2
                case NDIlib_FourCC_video_type_YV12: // 8-bit YUV 4:2:0
                case NDIlib_FourCC_video_type_I420: // 8-bit YUV 4:2:0
                case NDIlib_FourCC_video_type_NV12: // 8-bit YUV 4:2:0
                    // These are 4:2:2 or 4:2:0 formats. 
                    // On average, they use 2 bytes (16 bits) per pixel across the macroblock.
                    return 2; 

                case NDIlib_FourCC_video_type_BGRA: // 8-bit RGB with Alpha
                case NDIlib_FourCC_video_type_RGBA: // 8-bit RGB with Alpha
                    // 4 channels (Red, Green, Blue, Alpha) * 1 byte each
                    return 4; 

                case NDIlib_FourCC_video_type_BGRX: // 8-bit RGB (Padding)
                case NDIlib_FourCC_video_type_RGBX: // 8-bit RGB (Padding)
                    // 4 channels (Red, Green, Blue, Empty) * 1 byte each
                    return 4; 

                case NDIlib_FourCC_video_type_UYVA: // 8-bit YUV 4:2:2 + Alpha channel
                    // 2 bytes for YUV + 1 byte for Alpha split
                    return 3;

                case NDIlib_FourCC_video_type_P216: // 16-bit YUV 4:2:2 (High bit depth)
                    // 2 channels packed at 2 bytes (16-bits) per sample = 4 bytes per pixel
                    return 4;

                case NDIlib_FourCC_video_type_PA16: // 16-bit YUV 4:2:2 + 16-bit Alpha
                    return 6;

                default:
                    return 2; // Safe NDI default fallback
            }
        }

        std::string fourCCtoStr() {
            char fourcc_str[5];
            uint32_t fourcc = (uint32_t)fourCC;
            fourcc_str[0] = (fourcc >> 0)  & 0xFF;
            fourcc_str[1] = (fourcc >> 8)  & 0xFF;
            fourcc_str[2] = (fourcc >> 16) & 0xFF;
            fourcc_str[3] = (fourcc >> 24) & 0xFF;
            fourcc_str[4] = '\0';
            return std::string(fourcc_str);
        }

        bool getV210_video_frame(uint32_t source_num, uint8_t* frame_buffer, uint32_t frame_stride) {
            NDIlib_video_frame_v2_t video_frame;
            NDIlib_frame_type_e frame_type;
            frame_type = NDIlib_recv_capture_v3(pNDI_recv, &video_frame, nullptr, nullptr, 5);
            switch(frame_type)
            {
                case NDIlib_frame_type_error:
                    throw std::runtime_error("NDI source lost");
                case NDIlib_frame_type_status_change:
                    throw std::runtime_error("NDI source resolution or framerate are changed");

            }
            if (frame_type != NDIlib_frame_type_video) {
                return false;
            }
            switch(fourCC)
            {
                case NDIlib_FourCC_type_UYVY:
                    v210::UYVYtoV210(video_frame.p_data, frame_buffer, xres, yres, stride, frame_stride);
                    break;
                case NDIlib_FourCC_type_P216:
                    NDIlib_video_frame_v2_t video_frame_10bit;
                    NDIlib_util_P216_to_V210(&video_frame, &video_frame_10bit);
                    frame_buffer = video_frame.p_data;
                default:
                    throw std::runtime_error("Color format is not supported yet");
            }
            NDIlib_recv_free_video_v2(pNDI_recv, &video_frame);
            return true;
        }

    private:
        // receive
        NDIlib_find_instance_t pNDI_find = nullptr;
        NDIlib_recv_instance_t pNDI_recv = nullptr;
        NDIlib_FourCC_type_e fourCC;
        // owned copies so finder can be destroyed early
        std::vector<std::string>    cached_names;
        std::vector<std::string>    cached_urls;
        std::vector<NDIlib_source_t> cached_sources;
};
}