#include "../include/AsyncImageLoader.hpp"
#include "../include/DownloadUtils.hpp"
#include "../include/Program.hpp"
#include "../include/ImageUtils.hpp"
#include "../include/Scale.hpp"
#include "../include/Utils.hpp"
#include "../external/hash-library/sha256.h"

#include <unistd.h>
#include <sys/prctl.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <sys/time.h>
#include <fcntl.h>
#include <signal.h>
#include <malloc.h>
#include <assert.h>

#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include "../external/stb/stb_image_resize.h"

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wmissing-field-initializers"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "../external/stb/stb_image_write.h"
#pragma GCC diagnostic pop

namespace QuickMedia {
    static bool ffmpeg_image_to_png(const Path &thumbnail_path, const Path &destination_path) {
        const char *args[] = { "ffmpeg", "-y", "-v", "quiet", "-i", thumbnail_path.data.c_str(), "--", destination_path.data.c_str(), nullptr};
        return exec_program(args, nullptr, nullptr) == 0;
    }

    bool create_thumbnail(const Path &thumbnail_path, const Path &thumbnail_path_resized, mgl::vec2i resize_target_size, ContentType content_type) {
        Path input_path = thumbnail_path;

        if(content_type == ContentType::IMAGE_WEBP || content_type == ContentType::IMAGE_AVIF) {
            Path result_path_tmp = thumbnail_path_resized;
            result_path_tmp.append(".tmp.png");
            if(!ffmpeg_image_to_png(thumbnail_path, result_path_tmp))
                return false;
            input_path = std::move(result_path_tmp);
        }

        // Fork here because we want the memory allocated to be completely deallocated.
        // TODO: Find a way to do that without fork.
        pid_t parent_pid = getpid();
        pid_t pid = fork();
        if(pid == -1) {
            perror("Failed to fork");
            return false;
        } else if(pid == 0) { // child
            if(prctl(PR_SET_PDEATHSIG, SIGTERM) == -1) {
                perror("prctl(PR_SET_PDEATHSIG, SIGTERM) failed");
                _exit(127);
            }

            /* Test if the parent died before the above call to prctl */
            if(getppid() != parent_pid)
                _exit(127);

            mgl::Image image;
            if(!image.load_from_file(input_path.data.c_str()) || image.get_size().x == 0 || image.get_size().y == 0) {
                fprintf(stderr, "Failed to load %s\n", input_path.data.c_str());
                _exit(1);
            }

            Path result_path_tmp = thumbnail_path_resized;
            result_path_tmp.append(".tmp.png");

            if(image.get_size().x <= resize_target_size.x && image.get_size().y <= resize_target_size.y) {
                if(content_type == ContentType::IMAGE_WEBP || content_type == ContentType::IMAGE_AVIF) {
                    if(rename_atomic(input_path.data.c_str(), thumbnail_path_resized.data.c_str()) == 0)
                        _exit(0);
                    else
                        _exit(1);
                } else {
                    int res = symlink(thumbnail_path.data.c_str(), result_path_tmp.data.c_str());
                    if(res == -1 && errno != EEXIST) {
                        fprintf(stderr, "Failed to symlink %s to %s\n", result_path_tmp.data.c_str(), thumbnail_path.data.c_str());
                        _exit(1);
                    }
                }
            } else {
                mgl::vec2i clamped_size = clamp_to_size(image.get_size(), mgl::vec2i(resize_target_size.x, resize_target_size.y));
                unsigned char *output_pixels = new unsigned char[clamped_size.x * clamped_size.y * image.get_num_channels()];
                if(!stbir_resize_uint8(image.data(), image.get_size().x, image.get_size().y, 0, output_pixels, clamped_size.x, clamped_size.y, 0, image.get_num_channels())) {
                    fprintf(stderr, "Failed to resize %s\n", thumbnail_path_resized.data.c_str());
                    _exit(1);
                }

                if(image.get_num_channels() == 4) {
                    if(!stbi_write_png(result_path_tmp.data.c_str(), clamped_size.x, clamped_size.y, 4, output_pixels, 0)) {
                        fprintf(stderr, "Failed to save %s\n", thumbnail_path_resized.data.c_str());
                        _exit(1);
                    }
                } else {
                    if(!stbi_write_jpg(result_path_tmp.data.c_str(), clamped_size.x, clamped_size.y, image.get_num_channels(), output_pixels, 0)) {
                        fprintf(stderr, "Failed to save %s\n", thumbnail_path_resized.data.c_str());
                        _exit(1);
                    }
                }

                // TODO: Resize and save to dxt format which can be loaded faster, uses less memory during loading (since no conversion is needed to upload to gpu),
                // it uses less memory while in gpu (because its in dxt compressed format) and it gives better rendering performance because of compressed image = smaller = better cache utilization.
            }

            if(rename_atomic(result_path_tmp.data.c_str(), thumbnail_path_resized.data.c_str()) == 0)
                _exit(0);
            else
                _exit(1);
        }

        // parent

        // TODO: Do not wait here and instead check if finished in |load_thread|
        int status = 0;
        if(waitpid(pid, &status, 0) == -1) {
            perror("waitpid failed");
            return false;
        }

        if(!WIFEXITED(status))
            return false;

        int exit_status = WEXITSTATUS(status);
        if(exit_status != 0)
            return false;

        return true;
    }

    void AsyncImageLoader::load_create_thumbnail(const Path &thumbnail_path, const Path &thumbnail_path_resized, ThumbnailData *thumbnail_data, mgl::vec2i resize_target_size) {
        FileAnalyzer file_analyzer;
        if(!file_analyzer.load_file(thumbnail_path.data.c_str(), false)) {
            fprintf(stderr, "Failed to convert %s to a thumbnail\n", thumbnail_path.data.c_str());
            thumbnail_data->image = std::make_unique<mgl::Image>();
            thumbnail_data->loading_state = LoadingState::FAILED_TO_LOAD;
            return;
        }

        if(is_content_type_video(file_analyzer.get_content_type())) {
            if(file_analyzer.load_metadata() && video_get_middle_frame(file_analyzer, thumbnail_path_resized.data.c_str(), resize_target_size.x, resize_target_size.y)) {
                thumbnail_data->loading_state = LoadingState::READY_TO_LOAD;
            } else {
                fprintf(stderr, "Failed to get video frame of %s\n", thumbnail_path.data.c_str());
                thumbnail_data->image = std::make_unique<mgl::Image>();
                thumbnail_data->loading_state = LoadingState::FAILED_TO_LOAD;
            }
            return;
        }

        if(create_thumbnail(thumbnail_path, thumbnail_path_resized, resize_target_size, file_analyzer.get_content_type())) {
            thumbnail_data->loading_state = LoadingState::READY_TO_LOAD;
        } else {
            fprintf(stderr, "Failed to convert %s to a thumbnail\n", thumbnail_path.data.c_str());
            thumbnail_data->image = std::make_unique<mgl::Image>();
            thumbnail_data->loading_state = LoadingState::FAILED_TO_LOAD;
        }
    }

    AsyncImageLoader& AsyncImageLoader::get_instance() {
        static AsyncImageLoader *instance = nullptr;
        if(!instance)
            instance = new AsyncImageLoader();
        return *instance;
    }

    void AsyncImageLoader::process_thumbnail(ThumbnailLoadData &thumbnail_load_data) {
        Path thumbnail_path_resized = thumbnail_load_data.thumbnail_path;
        if(thumbnail_load_data.resize_target_size.x != 0 && thumbnail_load_data.resize_target_size.y != 0)
            thumbnail_path_resized.append("_" + std::to_string(thumbnail_load_data.resize_target_size.x) + "x" + std::to_string(thumbnail_load_data.resize_target_size.y));

        if(get_file_type(thumbnail_path_resized) == FileType::REGULAR) {
            fprintf(stderr, "Loaded %s from thumbnail cache\n", thumbnail_path_resized.data.c_str());
            thumbnail_load_data.thumbnail_data->image = std::make_unique<mgl::Image>();
            thumbnail_load_data.thumbnail_data->loading_state = LoadingState::READY_TO_LOAD;
            return;
        }

        Path thumbnail_original_path;
        if(thumbnail_load_data.local)
            thumbnail_original_path = thumbnail_load_data.path;
        else
            thumbnail_original_path = thumbnail_load_data.thumbnail_path;

        if(thumbnail_load_data.resize_target_size.x != 0 && thumbnail_load_data.resize_target_size.y != 0)
            load_create_thumbnail(thumbnail_original_path, thumbnail_path_resized, thumbnail_load_data.thumbnail_data.get(), thumbnail_load_data.resize_target_size);
        else
            thumbnail_load_data.thumbnail_data->loading_state = LoadingState::READY_TO_LOAD;
    }

    static void load_processed_thumbnail(ThumbnailLoadData &thumbnail_load_data) {
        thumbnail_load_data.thumbnail_data->image = std::make_unique<mgl::Image>();

        Path thumbnail_path_resized = thumbnail_load_data.thumbnail_path;
        if(thumbnail_load_data.resize_target_size.x != 0 && thumbnail_load_data.resize_target_size.y != 0)
            thumbnail_path_resized.append("_" + std::to_string(thumbnail_load_data.resize_target_size.x) + "x" + std::to_string(thumbnail_load_data.resize_target_size.y));

        if(get_file_type(thumbnail_path_resized) == FileType::REGULAR) {
            thumbnail_load_data.thumbnail_data->image->load_from_file(thumbnail_path_resized.data.c_str());
            thumbnail_load_data.thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
            return;
        }

        Path thumbnail_original_path;
        if(thumbnail_load_data.local)
            thumbnail_original_path = thumbnail_load_data.path;
        else
            thumbnail_original_path = thumbnail_load_data.thumbnail_path;

        if(thumbnail_load_data.resize_target_size.x != 0 && thumbnail_load_data.resize_target_size.y != 0) {
            thumbnail_load_data.thumbnail_data->image->load_from_file(thumbnail_path_resized.data.c_str());
            thumbnail_load_data.thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
        } else {
            thumbnail_load_data.thumbnail_data->image->load_from_file(thumbnail_original_path.data.c_str());
            thumbnail_load_data.thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
        }
    }

    AsyncImageLoader::AsyncImageLoader() {
        for(int i = 0; i < NUM_IMAGE_DOWNLOAD_PARALLEL; ++i) {
            downloads[i].read_program.pid = -1;
            downloads[i].read_program.read_fd = -1;
        }

        for(int i = 0; i < NUM_IMAGE_LOAD_PARALLEL; ++i) {
            load_threads[i] = AsyncTask<void>([this]() mutable {
                std::optional<ThumbnailLoadData> thumbnail_load_data_opt;
                while(image_thumbnail_create_queue.is_running()) {
                    thumbnail_load_data_opt = image_thumbnail_create_queue.pop_wait();
                    if(!thumbnail_load_data_opt)
                        break;

                    // TODO: Do this multithreaded because creating thumbnails is pretty slow single-threaded,
                    // especially video thumbnails.
                    process_thumbnail(thumbnail_load_data_opt.value());
                    if(thumbnail_load_data_opt.value().thumbnail_data->loading_state == LoadingState::READY_TO_LOAD)
                        load_processed_thumbnail(thumbnail_load_data_opt.value());

                    thumbnail_load_data_opt = std::nullopt;
                }
            });
        }
    }

    AsyncImageLoader::~AsyncImageLoader() {
        image_thumbnail_create_queue.close();
    }

    static bool download_file_async(const char *url, const char *save_filepath, ReadProgram *read_program) {
        std::string referer = "Referer: " + std::string(url);
        const char *args[] = {
            "curl", "-H", "Accept-Language: en-US,en;q=0.5", "-H", "Connection: keep-alive", "--compressed", "-g", "-s", "-L", "-f",
            "-H", "user-agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36",
            "-H", referer.c_str(),
            "-o", save_filepath,
            "--", url,
            nullptr
        };
        return exec_program_pipe(args, read_program) == 0;
    }

    bool AsyncImageLoader::load_thumbnail(const std::string &url, bool local, mgl::vec2i resize_target_size, std::shared_ptr<ThumbnailData> thumbnail_data, Path &thumbnail_path) {
        if(thumbnail_data->loading_state != LoadingState::NOT_LOADED)
            return true;

        if(url.empty()) {
            thumbnail_data->image = std::make_unique<mgl::Image>();
            thumbnail_data->loading_state = LoadingState::FAILED_TO_LOAD;
            return false;
        }
        
        if(local) {
            struct stat file_stat;
            memset(&file_stat, 0, sizeof(file_stat));
            if(stat(url.c_str(), &file_stat) != 0 || !S_ISREG(file_stat.st_mode)) {
                fprintf(stderr, "Failed to load thumbnail %s: no such file\n", url.c_str());
                thumbnail_data->image = std::make_unique<mgl::Image>();
                thumbnail_data->loading_state = LoadingState::FAILED_TO_LOAD;
                return false;
            }

            thumbnail_path.append("_" + std::to_string(file_stat.st_mtim.tv_sec));
            thumbnail_data->loading_state = LoadingState::LOADING;
            image_thumbnail_create_queue.push({ url, thumbnail_path, true, thumbnail_data, resize_target_size });
            return true;
        }

        if(get_file_type(thumbnail_path) == FileType::REGULAR) {
            thumbnail_data->loading_state = LoadingState::LOADING;
            image_thumbnail_create_queue.push({ url, thumbnail_path, false, thumbnail_data, resize_target_size });
            return true;
        }

        int free_index = get_free_load_index();
        if(free_index == -1)
            return false;

        std::lock_guard<std::mutex> lock(download_mutex);
        thumbnail_data->loading_state = LoadingState::LOADING;
        Path tmp_thumbnail_path = thumbnail_path;
        tmp_thumbnail_path.append(".tmp");
        if(!download_file_async(url.c_str(), tmp_thumbnail_path.data.c_str(), &downloads[free_index].read_program)) {
            fprintf(stderr, "Failed to start download of %s\n", url.c_str());
            return false;
        }

        downloads[free_index].download_start = get_boottime_milliseconds();
        downloads[free_index].thumbnail_path = thumbnail_path;
        downloads[free_index].thumbnail_data = thumbnail_data;
        downloads[free_index].resize_target_size = resize_target_size;
        downloads[free_index].url = url;
        return true;
    }

    std::shared_ptr<ThumbnailData> AsyncImageLoader::get_thumbnail(const std::string &url, bool local, mgl::vec2i resize_target_size) {
        // TODO: Instead of generating a new hash everytime to access thumbnail, cache the hash of the thumbnail url.
        // TODO: Cache this resize_url
        auto &thumbnail_data = thumbnails[url + "_" + std::to_string(resize_target_size.x) + "x" + std::to_string(resize_target_size.y)];
        if(!thumbnail_data)
            thumbnail_data = std::make_shared<ThumbnailData>();
        thumbnail_data->counter = counter;

        if(thumbnail_data->thumbnail_path.data.empty()) {
            SHA256 sha256;
            sha256.add(url.data(), url.size());
            thumbnail_data->thumbnail_path = get_cache_dir().join("thumbnails").join(sha256.getHash());
        }

        load_thumbnail(url, local, resize_target_size, thumbnail_data, thumbnail_data->thumbnail_path);
        return thumbnail_data;
    }

    void AsyncImageLoader::update() {
        for(int i = 0; i < NUM_IMAGE_DOWNLOAD_PARALLEL; ++i) {
            Download &download = downloads[i];
            if(download.read_program.pid == -1)
                continue;

            int status = 0;
            if(wait_program_non_blocking(download.read_program.pid, &status)) {
                Path tmp_thumbnail_path = download.thumbnail_path;
                tmp_thumbnail_path.append(".tmp");
                if(status == 0 && rename_atomic(tmp_thumbnail_path.data.c_str(), download.thumbnail_path.data.c_str()) == 0) {
                    fprintf(stderr, "Download duration for %s: %ld ms\n", download.url.c_str(), get_boottime_milliseconds() - download.download_start);
                    ThumbnailLoadData load_data = { std::move(download.url), std::move(download.thumbnail_path), false, download.thumbnail_data, download.resize_target_size };
                    image_thumbnail_create_queue.push(std::move(load_data));
                } else {
                    fprintf(stderr, "Thumbnail download failed for %s\n", download.url.c_str());
                    download.thumbnail_data->image = std::make_unique<mgl::Image>();
                    download.thumbnail_data->loading_state = LoadingState::FAILED_TO_LOAD;
                }
                reset_download(download);
            }
        }

        bool loaded_textures_changed = false;
        for(auto it = thumbnails.begin(); it != thumbnails.end();) {
            if(it->second->counter != counter) {
                {
                    for(int i = 0; i < NUM_IMAGE_DOWNLOAD_PARALLEL; ++i) {
                        Download &download = downloads[i];
                        if(download.read_program.pid == -1)
                            continue;

                        if(download.thumbnail_data.get() == it->second.get()) {
                            reset_download(download);
                            break;
                        }
                    }
                }

                image_thumbnail_create_queue.erase_if([&it](ThumbnailLoadData &load_data) {
                    return load_data.thumbnail_data.get() == it->second.get();
                });

                it = thumbnails.erase(it);
                loaded_textures_changed = true;
            } else {
                ++it;
            }
        }

        ++counter;
        if(loaded_textures_changed)
            malloc_trim(0);
    }

    int AsyncImageLoader::get_free_load_index() const {
        for(int i = 0; i < NUM_IMAGE_DOWNLOAD_PARALLEL; ++i) {
            if(downloads[i].read_program.pid == -1)
                return i;
        }
        return -1;
    }

    void AsyncImageLoader::reset_download(Download &download) {
        std::lock_guard<std::mutex> lock(download_mutex);
        if(download.read_program.pid != -1) {
            kill(download.read_program.pid, SIGKILL);
            int status;
            waitpid(download.read_program.pid, &status, 0);
            download.read_program.pid = -1;
        }
        if(download.read_program.read_fd != -1) {
            close(download.read_program.read_fd);
            download.read_program.read_fd = -1;
        }
        download.thumbnail_path.data.clear();
        download.thumbnail_data = nullptr;
        download.url.clear();
    }
}