path: root/src/AsyncImageLoader.cpp

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

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

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

namespace QuickMedia {
    static bool webp_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, sf::Vector2i resize_target_size, ContentType content_type) {
        Path input_path = thumbnail_path;

        if(content_type == ContentType::IMAGE_WEBP) {
            Path result_path_tmp = thumbnail_path_resized;
            result_path_tmp.append(".tmp.png");
            if(!webp_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);

            sf::Image image;
            if(!image.loadFromFile(input_path.data) || image.getSize().x == 0 || image.getSize().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.getSize().x <= (unsigned int)resize_target_size.x && image.getSize().y <= (unsigned int)resize_target_size.y) {
                int res = symlink(input_path.data.c_str(), result_path_tmp.data.c_str());
                if(res == -1 && errno != EEXIST) {
                    fprintf(stderr, "Failed to save %s\n", thumbnail_path_resized.data.c_str());
                    _exit(1);
                }
            } else {
                sf::Vector2u clamped_size = clamp_to_size(image.getSize(), sf::Vector2u(resize_target_size.x, resize_target_size.y));
                unsigned char *output_pixels = new unsigned char[clamped_size.x * clamped_size.y * 4];
                stbir_resize_uint8(image.getPixelsPtr(), image.getSize().x, image.getSize().y, 0, output_pixels, clamped_size.x, clamped_size.y, 0, 4);

                // TODO: Remove this and use stb write to remove this unecessary extra copy of the data and write the data directly to file after converting it to png
                sf::Image destination_image;
                destination_image.create(clamped_size.x, clamped_size.y, output_pixels);
                if(!destination_image.saveToFile(result_path_tmp.data)) {
                    fprintf(stderr, "Failed to save %s\n", thumbnail_path_resized.data.c_str());
                    _exit(1);
                }
            }

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

        // parent

        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;
    }

    // Create thumbnail and load it. On failure load the original image
    static void create_thumbnail(const Path &thumbnail_path, const Path &thumbnail_path_resized, ThumbnailData *thumbnail_data, sf::Vector2i 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, using the original image\n", thumbnail_path.data.c_str());
            thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
            return;
        }

        if(is_content_type_video(file_analyzer.get_content_type())) {
            if(video_get_first_frame(thumbnail_path.data.c_str(), thumbnail_path_resized.data.c_str(), resize_target_size.x, resize_target_size.y))
                load_image_from_file(*thumbnail_data->image, thumbnail_path_resized.data);
            thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
            return;
        }

        if(create_thumbnail(thumbnail_path, thumbnail_path_resized, resize_target_size, file_analyzer.get_content_type())) {
            load_image_from_file(*thumbnail_data->image, thumbnail_path_resized.data);
        } else {
            load_image_from_file(*thumbnail_data->image, thumbnail_path.data);
            fprintf(stderr, "Failed to convert %s to a thumbnail, using the original image\n", thumbnail_path.data.c_str());
        }
        thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
    }

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

    AsyncImageLoader::AsyncImageLoader() {
        for(int i = 0; i < NUM_IMAGE_LOAD_THREADS; ++i) {
            loading_image[i] = false;
        }

        load_image_thread = AsyncTask<void>([this]() mutable {
            std::optional<ThumbnailLoadData> thumbnail_load_data_opt;
            while(true) {
                thumbnail_load_data_opt = image_load_queue.pop_wait();
                if(!thumbnail_load_data_opt)
                    break;

                ThumbnailLoadData &thumbnail_load_data = thumbnail_load_data_opt.value();
                thumbnail_load_data.thumbnail_data->image = std::make_unique<sf::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) {
                    load_image_from_file(*thumbnail_load_data.thumbnail_data->image, thumbnail_path_resized.data);
                    thumbnail_load_data.thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
                    fprintf(stderr, "Loaded %s from thumbnail cache\n", thumbnail_path_resized.data.c_str());
                    continue;
                }

                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)
                    create_thumbnail(thumbnail_original_path, thumbnail_path_resized, thumbnail_load_data.thumbnail_data.get(), thumbnail_load_data.resize_target_size);
                else
                    load_image_from_file(*thumbnail_load_data.thumbnail_data->image, thumbnail_original_path.data);

                thumbnail_load_data.thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
            }
        });
    }

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

    void AsyncImageLoader::load_thumbnail(const std::string &url, bool local, sf::Vector2i resize_target_size, std::shared_ptr<ThumbnailData> thumbnail_data) {
        if(thumbnail_data->loading_state != LoadingState::NOT_LOADED)
            return;

        if(url.empty()) {
            thumbnail_data->image = std::make_unique<sf::Image>();
            thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
            return;
        }
        
        SHA256 sha256;
        sha256.add(url.data(), url.size());
        Path thumbnail_path = get_cache_dir().join("thumbnails").join(sha256.getHash());
        if(local) {
            struct stat file_stat;
            if(stat(url.c_str(), &file_stat) != 0 || !S_ISREG(file_stat.st_mode)) {
                thumbnail_data->image = std::make_unique<sf::Image>();
                thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
                return;
            }

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

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

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

        loading_image[free_index] = true;
        thumbnail_data->loading_state = LoadingState::LOADING;

        // TODO: Keep the thread running and use conditional variable instead to sleep until a new image should be loaded. Same in ImageViewer.
        download_image_thread[free_index] = AsyncTask<void>([this, free_index, thumbnail_path, url, resize_target_size, thumbnail_data]() mutable {
            thumbnail_data->image = std::make_unique<sf::Image>();

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

            if(get_file_type(thumbnail_path_resized) == FileType::REGULAR) {
                load_image_from_file(*thumbnail_data->image, thumbnail_path_resized.data);
                thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
                fprintf(stderr, "Loaded %s from thumbnail cache\n", thumbnail_path_resized.data.c_str());
                return;
            }

            if(get_file_type(thumbnail_path.data) == FileType::FILE_NOT_FOUND && download_to_file(url, thumbnail_path.data, {}, true) != DownloadResult::OK) {
                thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
                loading_image[free_index] = false;
                return;
            }

            if(resize_target_size.x != 0 && resize_target_size.y != 0)
                create_thumbnail(thumbnail_path, thumbnail_path_resized, thumbnail_data.get(), resize_target_size);
            else
                load_image_from_file(*thumbnail_data->image, thumbnail_path.data);

            thumbnail_data->loading_state = LoadingState::FINISHED_LOADING;
            loading_image[free_index] = false;
            return;
        });
    }

    std::shared_ptr<ThumbnailData> AsyncImageLoader::get_thumbnail(const std::string &url, bool local, sf::Vector2i resize_target_size) {
        // TODO: Instead of generating a new hash everytime to access thumbnail, cache the hash of the thumbnail url
        auto &thumbnail_data = thumbnails[url];
        if(!thumbnail_data)
            thumbnail_data = std::make_shared<ThumbnailData>();
        thumbnail_data->counter = counter;
        load_thumbnail(url, local, resize_target_size, thumbnail_data);
        return thumbnail_data;
    }

    void AsyncImageLoader::update() {
        bool loaded_textures_changed = false;
        for(auto it = thumbnails.begin(); it != thumbnails.end();) {
            if(it->second->counter != counter) {
                image_load_queue.erase_if([&it](ThumbnailLoadData &load_data) {
                    return load_data.path.data == it->first;
                });
                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_LOAD_THREADS; ++i) {
            if(!loading_image[i])
                return i;
        }
        return -1;
    }
}