#include "module_gvar_image_decoder.h"
#include "logger.h"
#include "common/codings/differential/nrzs.h"
#include "imgui/imgui.h"
#include "gvar_deframer.h"
#include "gvar_derand.h"
#include "imgui/imgui_image.h"
#include <filesystem>
#include "gvar_headers.h"
#include "common/utils.h"
#include "resources.h"
#include "common/image/hue_saturation.h"
#include "common/image/brightness_contrast.h"
#include "common/thread_priority.h"
#include "common/image/io.h"

#define FRAME_SIZE 32786

// Return filesize
uint64_t getFilesize(std::string filepath);

namespace goes
{
    namespace gvar
    {
        std::string GVARImageDecoderModule::getGvarFilename(int sat_number, std::tm *timeReadable, std::string channel)
        {
            std::string utc_filename = "G" + std::to_string(sat_number) + "_" + channel + "_" +                                                                     // Satellite name and channel
                                       std::to_string(timeReadable->tm_year + 1900) +                                                                               // Year yyyy
                                       (timeReadable->tm_mon + 1 > 9 ? std::to_string(timeReadable->tm_mon + 1) : "0" + std::to_string(timeReadable->tm_mon + 1)) + // Month MM
                                       (timeReadable->tm_mday > 9 ? std::to_string(timeReadable->tm_mday) : "0" + std::to_string(timeReadable->tm_mday)) + "T" +    // Day dd
                                       (timeReadable->tm_hour > 9 ? std::to_string(timeReadable->tm_hour) : "0" + std::to_string(timeReadable->tm_hour)) +          // Hour HH
                                       (timeReadable->tm_min > 9 ? std::to_string(timeReadable->tm_min) : "0" + std::to_string(timeReadable->tm_min)) +             // Minutes mm
                                       (timeReadable->tm_sec > 9 ? std::to_string(timeReadable->tm_sec) : "0" + std::to_string(timeReadable->tm_sec)) + "Z";        // Seconds ss
            return utc_filename;
        }

        void GVARImageDecoderModule::writeSounder()
        {
            const time_t timevalue = time(0);
            std::tm *timeReadable = gmtime(&timevalue);
            std::string timestamp = std::to_string(timeReadable->tm_year + 1900) + "-" +
                                    (timeReadable->tm_mon + 1 > 9 ? std::to_string(timeReadable->tm_mon + 1) : "0" + std::to_string(timeReadable->tm_mon + 1)) + "-" +
                                    (timeReadable->tm_mday > 9 ? std::to_string(timeReadable->tm_mday) : "0" + std::to_string(timeReadable->tm_mday)) + "_" +
                                    (timeReadable->tm_hour > 9 ? std::to_string(timeReadable->tm_hour) : "0" + std::to_string(timeReadable->tm_hour)) + "-" +
                                    (timeReadable->tm_min > 9 ? std::to_string(timeReadable->tm_min) : "0" + std::to_string(timeReadable->tm_min));

            logger->info("New sounder scan starting, saving at " + timestamp + "...");

            std::filesystem::create_directories(directory_snd + "/" + timestamp);
            std::string directory_d = directory_snd + "/" + timestamp;

            for (int i = 0; i < 19; i++)
            {
                auto img = sounderReader.getImage(i);
                image::save_img(img, directory_d + "/Sounder_" + std::to_string(i + 1));
            }

            sounderReader.clear();
        }

        void GVARImageDecoderModule::writeImages(GVARImages &images, std::string directory)
        {
            const time_t timevalue = time(0);
            std::tm *timeReadable = gmtime(&timevalue);
            std::string timestamp = std::to_string(timeReadable->tm_year + 1900) + "-" +
                                    (timeReadable->tm_mon + 1 > 9 ? std::to_string(timeReadable->tm_mon + 1) : "0" + std::to_string(timeReadable->tm_mon + 1)) + "-" +
                                    (timeReadable->tm_mday > 9 ? std::to_string(timeReadable->tm_mday) : "0" + std::to_string(timeReadable->tm_mday)) + "_" +
                                    (timeReadable->tm_hour > 9 ? std::to_string(timeReadable->tm_hour) : "0" + std::to_string(timeReadable->tm_hour)) + "-" +
                                    (timeReadable->tm_min > 9 ? std::to_string(timeReadable->tm_min) : "0" + std::to_string(timeReadable->tm_min));

            // Get stats. This is done over a lot of data to allow decoding at low SNR
            int sat_number = images.sat_number;
            int vis_width = images.vis_width;

            std::string dir_name = "GOES-" + std::to_string(sat_number) + "/" + timestamp;
            logger->info("Full disk finished, saving at " + dir_name + "...");

            std::filesystem::create_directories(directory + "/" + dir_name);

            std::string disk_folder = directory + "/" + dir_name;

            logger->info("Resizing...");
            images.image1.resize(images.image1.width(), images.image1.height() * 1.75);
            images.image2.resize(images.image2.width(), images.image2.height() * 1.75);
            images.image3.resize(images.image3.width(), images.image3.height() * 1.75);
            images.image4.resize(images.image4.width(), images.image4.height() * 1.75);
            images.image5.resize(images.image5.width(), images.image5.height() * 1.75);

            // logger->trace("VIS 1 size before " + std::to_string(images.image5.width()) + "x" + std::to_string(images.image5.height()));
            // logger->trace("IR size before " + std::to_string(images.image1.width()) + "x" + std::to_string(images.image1.height()));

            // VIS-1 height
            int vis_height = images.image5.height();

            if (vis_width == 13216) // Some partial scan
                vis_height = 9500;
            else if (vis_width == 11416) // Some partial scan
                vis_height = 6895;
            else if (vis_width == 8396) // Some partial scan
                vis_height = 4600;
            else if (vis_width == 11012) // Some partial scan
                vis_height = 7456;
            else if (vis_width == 4976) // Some partial scan
                vis_height = 4194;
            else if (vis_width == 20836) // Full disk
                vis_height = 18956;
            else if (vis_width == 20824) // Full disk
                vis_height = 18956;

            // IR height
            int ir1_width = vis_width / 4;
            int ir1_height = vis_height / 4;

            logger->info("Cropping to transmited size...");
            logger->debug("VIS 1 size " + std::to_string(vis_width) + "x" + std::to_string(vis_height));
            images.image5.crop(0, 0, vis_width, vis_height);
            logger->debug("IR size " + std::to_string(ir1_width) + "x" + std::to_string(ir1_height));
            images.image1.crop(0, 0, ir1_width, ir1_height);
            images.image2.crop(0, 0, ir1_width, ir1_height);
            images.image3.crop(0, 0, ir1_width, ir1_height);
            images.image4.crop(0, 0, ir1_width, ir1_height);

            image::save_img(images.image5, std::string(disk_folder + "/" + getGvarFilename(sat_number, timeReadable, "1")).c_str(), false);
            image::save_img(images.image1, std::string(disk_folder + "/" + getGvarFilename(sat_number, timeReadable, "2")).c_str());
            image::save_img(images.image2, std::string(disk_folder + "/" + getGvarFilename(sat_number, timeReadable, "3")).c_str());
            image::save_img(images.image3, std::string(disk_folder + "/" + getGvarFilename(sat_number, timeReadable, "4")).c_str());
            image::save_img(images.image4, std::string(disk_folder + "/" + getGvarFilename(sat_number, timeReadable, "5")).c_str());

            // Let plugins do something
            satdump::eventBus->fire_event<events::GVARSaveChannelImagesEvent>({images, timeReadable, timevalue, disk_folder});

            // We are done with all channels but 1 and 4. Clear others to free up memory!
            images.image1.clear();
            images.image2.clear();
            images.image4.clear();

            // If we can, generate false color
            if (resources::resourceExists("goes/gvar/lut.png"))
            {
                logger->trace("Scale Ch1 to 8-bits...");
                image::Image channel1(8, images.image5.width(), images.image5.height(), 1);
                for (size_t i = 0; i < channel1.width() * channel1.height(); i++)
                    channel1.set(i, images.image5.get(i) / 255);
                images.image5.clear(); // We're done with Ch1. Free up memory

                logger->trace("Scale Ch4 to 8-bits...");
                image::Image channel4(8, images.image3.width(), images.image3.height(), 1);
                for (size_t i = 0; i < channel4.width() * channel4.height(); i++)
                    channel4.set(i, images.image3.get(i) / 255);
                images.image3.clear(); // We're done with Ch4. Free up memory

                logger->trace("Resize images...");
                channel4.resize(channel1.width(), channel1.height());

                logger->trace("Loading LUT...");
                image::Image lutImage;
                image::load_png(lutImage, resources::getResourcePath("goes/gvar/lut.png").c_str());
                lutImage.resize(256, 256);

                logger->trace("Loading correction curve...");
                image::Image curveImage;
                image::load_png(curveImage, resources::getResourcePath("goes/gvar/curve_goesn.png").c_str());

                image::Image compoImage(8, channel1.width(), channel1.height(), 3);

                logger->trace("Applying LUT...");
                for (size_t i = 0; i < channel1.width() * channel1.height(); i++)
                {
                    uint8_t x = 255 - curveImage.get(channel1.get(i)) / 1.5;
                    uint8_t y = channel4.get(i);

                    for (int c = 0; c < 3; c++)
                        compoImage.set(c * compoImage.width() * compoImage.height() + i, lutImage.get(c * lutImage.width() * lutImage.height() + x * lutImage.width() + y));
                }

                logger->trace("Contrast correction...");
                image::brightness_contrast(compoImage, -10.0f / 127.0f, 24.0f / 127.0f);

                logger->trace("Hue shift...");
                image::HueSaturation hueTuning;
                hueTuning.hue[image::HUE_RANGE_MAGENTA] = 133.0 / 180.0;
                hueTuning.overlap = 100.0 / 100.0;
                image::hue_saturation(compoImage, hueTuning);
                image::save_img(compoImage, std::string(disk_folder + "/" + getGvarFilename(sat_number, timeReadable, "FC")).c_str());

                // Let plugins do something
                satdump::eventBus->fire_event<events::GVARSaveFCImageEvent>({compoImage, images.sat_number, timeReadable, timevalue, disk_folder});
            }
            else
            {
                logger->warn("goes/gvar/lut.png LUT is missing! False Color will not be generated");
            }
        }

        GVARImageDecoderModule::GVARImageDecoderModule(std::string input_file, std::string output_file_hint, nlohmann::json parameters) : ProcessingModule(input_file, output_file_hint, parameters)
        {
            frame = new uint8_t[FRAME_SIZE];
            isImageInProgress = false;
            isSavingInProgress = false;
            approx_progess = 0;

            // Reset readers
            infraredImageReader1.startNewFullDisk();
            infraredImageReader2.startNewFullDisk();
            visibleImageReader.startNewFullDisk();

            nonEndCount = 0;
            endCount = 0;
        }

        std::vector<ModuleDataType> GVARImageDecoderModule::getInputTypes()
        {
            return {DATA_FILE, DATA_STREAM};
        }

        std::vector<ModuleDataType> GVARImageDecoderModule::getOutputTypes()
        {
            return {DATA_FILE};
        }

        GVARImageDecoderModule::~GVARImageDecoderModule()
        {
            delete[] frame;

            if (textureID != 0)
            {
                delete[] textureBuffer;
                // deleteImageTexture(textureID);
            }
        }

        void GVARImageDecoderModule::writeImagesThread()
        {
            logger->info("Started saving thread...");
            while (writeImagesAync)
            {
                imageVectorMutex.lock();
                int imagesCount = imagesVector.size();
                if (imagesCount > 0)
                {
                    writeImages(imagesVector[0], directory);
                    imagesVector.erase(imagesVector.begin(), imagesVector.begin() + 1);
                }
                imageVectorMutex.unlock();
                std::this_thread::sleep_for(std::chrono::seconds(1));
            }
        }

        void GVARImageDecoderModule::process()
        {
            if (input_data_type == DATA_FILE)
                filesize = getFilesize(d_input_file);
            else
                filesize = 0;
            if (input_data_type == DATA_FILE)
                data_in = std::ifstream(d_input_file, std::ios::binary);

            if (input_data_type != DATA_FILE)
                writeImagesAync = true;

            // Start thread
            if (writeImagesAync)
            {
                imageSavingThread = std::thread(&GVARImageDecoderModule::writeImagesThread, this);
                setLowestThreadPriority(imageSavingThread); // Low priority to avoid sampledrop
            }

            directory = d_output_file_hint.substr(0, d_output_file_hint.rfind('/')) + "/IMAGE";
            directory_snd = d_output_file_hint.substr(0, d_output_file_hint.rfind('/')) + "/SOUNDER";

            if (!std::filesystem::exists(directory))
                std::filesystem::create_directory(directory);

            if (!std::filesystem::exists(directory_snd))
                std::filesystem::create_directory(directory_snd);

            logger->info("Using input frames " + d_input_file);
            logger->info("Decoding to " + directory);

            time_t lastTime = 0;

            int last_val = 0;

            while (input_data_type == DATA_FILE ? !data_in.eof() : input_active.load())
            {
                // Read a buffer
                if (input_data_type == DATA_FILE)
                    data_in.read((char *)frame, FRAME_SIZE);
                else
                    input_fifo->read((uint8_t *)frame, FRAME_SIZE);

                // Parse main header
                PrimaryBlockHeader block_header = *((PrimaryBlockHeader *)&frame[8]);

                if (block_header.block_id == 11)
                {
                    uint8_t *sad_header = &frame[8 + 30 * 3];
                    int dataid = sad_header[2];

                    if (dataid == 35)
                    {
                        int x_pos = frame[6120] << 8 | frame[6121];
                        if (sad_header[3])
                        {
                            if (x_pos != last_val)
                                writeSounder();
                            last_val = x_pos;
                        }
                        sounderReader.pushFrame(frame, 0);
                    }
                }

                // Is this imagery? Blocks 1 to 10 are imagery
                if (block_header.block_id >= 1 && block_header.block_id <= 10)
                {
                    // This is imagery, so we can parse the line information header
                    LineDocumentationHeader line_header(&frame[8 + 30 * 3]);

                    // SCID Stats
                    scid_stats.push_back(line_header.sc_id);

                    // Internal line counter should NEVER be over 1974. Discard frame if it is
                    // Though we know the max in normal operations is 1354 for a full disk....
                    // So we use that
                    if (line_header.relative_scan_count > 1354)
                        continue;

                    // Is this VIS Channel 1?
                    if (block_header.block_id >= 3 && block_header.block_id <= 10)
                    {
                        // Push width stats
                        vis_width_stats.push_back(line_header.pixel_count);

                        // Push into decoder
                        visibleImageReader.pushFrame(frame, block_header.block_id, line_header.relative_scan_count);

                        // Detect full disk end
                        if (line_header.relative_scan_count < 2)
                        {
                            nonEndCount = 0;
                            endCount += 2;

                            if (endCount > 6)
                            {
                                logger->info("Image start detected!");

                                if (isImageInProgress)
                                {
                                    if (writeImagesAync)
                                    {
                                        logger->debug("Saving Async...");
                                        isImageInProgress = false;
                                        isSavingInProgress = true;
                                        imageVectorMutex.lock();
                                        imagesVector.push_back({infraredImageReader1.getImage1(),
                                                                infraredImageReader1.getImage2(),
                                                                infraredImageReader2.getImage1(),
                                                                infraredImageReader2.getImage2(),
                                                                visibleImageReader.getImage(),
                                                                most_common(scid_stats.begin(), scid_stats.end(), 0),
                                                                most_common(vis_width_stats.begin(), vis_width_stats.end(), 0)});
                                        imageVectorMutex.unlock();
                                        isSavingInProgress = false;
                                    }
                                    else
                                    {
                                        logger->debug("Saving...");
                                        isImageInProgress = false;
                                        isSavingInProgress = true;
                                        GVARImages images = {infraredImageReader1.getImage1(),
                                                             infraredImageReader1.getImage2(),
                                                             infraredImageReader2.getImage1(),
                                                             infraredImageReader2.getImage2(),
                                                             visibleImageReader.getImage(),
                                                             most_common(scid_stats.begin(), scid_stats.end(), 0),
                                                             most_common(vis_width_stats.begin(), vis_width_stats.end(), 0)};
                                        writeImages(images, directory);
                                        isSavingInProgress = false;
                                    }

                                    scid_stats.clear();
                                    vis_width_stats.clear();
                                    ir_width_stats.clear();

                                    // Reset readers
                                    infraredImageReader1.startNewFullDisk();
                                    infraredImageReader2.startNewFullDisk();
                                    visibleImageReader.startNewFullDisk();
                                }

                                endCount = 0;
                            }
                        }
                        else
                        {
                            if (endCount > 0)
                                endCount--;

                            nonEndCount++;

                            if (nonEndCount > 100 && !isImageInProgress)
                                isImageInProgress = true;
                        }
                    }
                    // Is this IR?
                    else if (block_header.block_id == 1 || block_header.block_id == 2)
                    {
                        // Easy way of showing an approximate progress percentage
                        approx_progess = round(((float)line_header.relative_scan_count / 1353.0f) * 1000.0f) / 10.0f;

                        // Push frame size stats
                        ir_width_stats.push_back(line_header.word_count);

                        // Get current stats
                        int current_words = most_common(ir_width_stats.begin(), ir_width_stats.end(), 0);

                        // Safeguard
                        if (current_words > 6565)
                            current_words = 6530; // Default to fulldisk size

                        // Is this IR Channel 1-2?
                        if (block_header.block_id == 1)
                            infraredImageReader1.pushFrame(&frame[8 + 30 * 3], line_header.relative_scan_count, current_words);
                        else if (block_header.block_id == 2)
                            infraredImageReader2.pushFrame(&frame[8 + 30 * 3], line_header.relative_scan_count, current_words);
                    }
                }

                if (input_data_type == DATA_FILE)
                    progress = data_in.tellg();

                // Update module stats
                module_stats["full_disk_progress"] = approx_progess;

                if (time(NULL) % 10 == 0 && lastTime != time(NULL))
                {
                    lastTime = time(NULL);
                    logger->info("Progress " + std::to_string(round(((double)progress / (double)filesize) * 1000.0) / 10.0) +
                                 "%%, Full Disk Progress : " + std::to_string(round(((float)approx_progess / 100.0f) * 1000.0f) / 10.0f) + "%%");
                }
            }

            if (input_data_type == DATA_FILE)
                data_in.close();

            writeSounder();

            if (writeImagesAync)
            {
                logger->info("Exit async thread...");
                writeImagesAync = false;
                if (imageSavingThread.joinable())
                    imageSavingThread.join();
            }

            logger->info("Dump remaining data...");
            if (isImageInProgress)
            {
                isImageInProgress = false;
                isSavingInProgress = true;
                // Backup images
                GVARImages images = {infraredImageReader1.getImage1(),
                                     infraredImageReader1.getImage2(),
                                     infraredImageReader2.getImage1(),
                                     infraredImageReader2.getImage2(),
                                     visibleImageReader.getImage(),
                                     most_common(scid_stats.begin(), scid_stats.end(), 0),
                                     most_common(vis_width_stats.begin(), vis_width_stats.end(), 0)};
                // Write those
                writeImages(images, directory);
            }
        }

        void GVARImageDecoderModule::drawUI(bool window)
        {
            if (textureID == 0)
            {
                textureID = makeImageTexture();
                textureBuffer = new uint32_t[1354 * 2 * 5236];
                memset(textureBuffer, 0, sizeof(uint32_t) * 1354 * 2 * 5236);
            }

            ImGui::Begin("GVAR Image Decoder", NULL, window ? 0 : NOWINDOW_FLAGS);

            // This is outer crap...
            ImGui::BeginGroup();
            {
                ushort_to_rgba(infraredImageReader1.imageBuffer1, textureBuffer, 5236 * 1354 * 2);
                updateImageTexture(textureID, textureBuffer, 5236, 1354 * 2);
                ImGui::Image((void *)(intptr_t)textureID, {200 * ui_scale, 200 * ui_scale});
            }
            ImGui::EndGroup();

            ImGui::SameLine();

            ImGui::BeginGroup();
            {
                ImGui::Button("Full Disk Progress", {200 * ui_scale, 20 * ui_scale});
                ImGui::ProgressBar((float)approx_progess / 100.0f, ImVec2(200 * ui_scale, 20 * ui_scale));
                ImGui::Text("State : ");
                ImGui::SameLine();
                if (isSavingInProgress)
                    ImGui::TextColored(style::theme.green, "Writing images...");
                else if (isImageInProgress)
                    ImGui::TextColored(style::theme.orange, "Receiving...");
                else
                    ImGui::TextColored(style::theme.red, "IDLE");
            }
            ImGui::EndGroup();

            if (!streamingInput)
                ImGui::ProgressBar((double)progress / (double)filesize, ImVec2(ImGui::GetContentRegionAvail().x, 20 * ui_scale));

            ImGui::End();
        }

        std::string GVARImageDecoderModule::getID()
        {
            return "goes_gvar_image_decoder";
        }

        std::vector<std::string> GVARImageDecoderModule::getParameters()
        {
            return {};
        }

        std::shared_ptr<ProcessingModule> GVARImageDecoderModule::getInstance(std::string input_file, std::string output_file_hint, nlohmann::json parameters)
        {
            return std::make_shared<GVARImageDecoderModule>(input_file, output_file_hint, parameters);
        }
    } // namespace elektro
}