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 "../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/sendfile.h>
#include <sys/time.h>
#include <fcntl.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, bool symlink_if_no_resize) {
        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) {
                if(content_type == ContentType::IMAGE_WEBP) {
                    if(rename_atomic(input_path.data.c_str(), thumbnail_path_resized.data.c_str()) == 0)
                        _exit(0);
                    else
                        _exit(1);
                } else if(symlink_if_no_resize) {
                    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", thumbnail_path_resized.data.c_str(), thumbnail_path.data.c_str());
                        _exit(1);
                    }
                } else {
                    // TODO: When mac is supported (or other OS than linux), then fix this for them. Mac for example needs fcopyfile instead of sendfile
                    int input_file = open(thumbnail_path.data.c_str(), O_RDONLY);
                    if(input_file == -1) {
                        fprintf(stderr, "Failed to save %s\n", thumbnail_path_resized.data.c_str());
                        _exit(1);
                    }

                    int output_file = creat(result_path_tmp.data.c_str(), 0660);
                    if(output_file == -1) {
                        fprintf(stderr, "Failed to save %s\n", thumbnail_path_resized.data.c_str());
                        _exit(1);
                    }

                    off_t bytes_copied = 0;
                    struct stat file_stat;
                    memset(&file_stat, 0, sizeof(file_stat));
                    if(fstat(input_file, &file_stat) == -1) {
                        fprintf(stderr, "Failed to save %s\n", thumbnail_path_resized.data.c_str());
                        _exit(1);
                    }

                    // No need to retry, small files
                    if(sendfile(output_file, input_file, &bytes_copied, file_stat.st_size) == -1) {
                         fprintf(stderr, "Failed to save %s\n", thumbnail_path_resized.data.c_str());
                        _exit(1);
                    }

                    close(input_file);
                    close(output_file);
                }
            } 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

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

    // 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(), true)) {
            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;
    }

    static void process_thumbnail(ThumbnailLoadData &thumbnail_load_data) {
        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());
            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)
            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() {
        for(int i = 0; i < NUM_IMAGE_LOAD_PARALLEL; ++i) {
            downloads[i].read_program.pid = -1;
            downloads[i].read_program.read_fd = -1;
        }

        load_thread = AsyncTask<void>([this]() mutable {
            std::optional<ThumbnailLoadData> thumbnail_load_data_opt;
            while(image_load_queue.is_running()) {
                thumbnail_load_data_opt = image_load_queue.pop_if_available();
                if(thumbnail_load_data_opt)
                    process_thumbnail(thumbnail_load_data_opt.value());

                for(int i = 0; i < NUM_IMAGE_LOAD_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 };
                            process_thumbnail(load_data);
                        } else {
                            fprintf(stderr, "Thumbnail download failed for %s\n", download.url.c_str());
                        }

                        std::lock_guard<std::mutex> lock(download_mutex);
                        close(download.read_program.read_fd);
                        download.read_program.pid = -1;
                        download.read_program.read_fd = -1;
                        download.thumbnail_path.data.clear();
                        download.url.c_str();
                        download.thumbnail_data = nullptr;
                    }
                }

                std::this_thread::sleep_for(std::chrono::milliseconds(2));
            }
        });
    }

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

    static bool download_file_async(const char *url, const char *save_filepath, ReadProgram *read_program) {
        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",
            "-o", save_filepath,
            "--", url,
            nullptr
        };
        return exec_program_pipe(args, read_program) == 0;
    }

    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;
            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<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, std::move(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, std::move(thumbnail_path), false, thumbnail_data, resize_target_size });
            return;
        }

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

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

        downloads[free_index].download_start = get_boottime_milliseconds();
        downloads[free_index].thumbnail_path = std::move(thumbnail_path);
        downloads[free_index].thumbnail_data = thumbnail_data;
        downloads[free_index].resize_target_size = resize_target_size;
        downloads[free_index].url = url;
    }

    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_PARALLEL; ++i) {
            if(downloads[i].read_program.pid == -1)
                return i;
        }
        return -1;
    }
}