path: root/src/Downloader.cpp

#include "../include/Downloader.hpp"
#include "../include/Storage.hpp"
#include "../include/NetUtils.hpp"
#include "../include/Notification.hpp"
#include <unistd.h>
#include <signal.h>

namespace QuickMedia {
    static bool youtube_url_is_live_stream(const std::string &url) {
        return url.find("yt_live_broadcast") != std::string::npos || url.find("manifest/") != std::string::npos;
    }

    CurlDownloader::CurlDownloader(const std::string &url, const std::string &output_filepath, int64_t content_length) : Downloader(url, output_filepath), content_length(content_length) {
        output_filepath_tmp = output_filepath;
        output_filepath_tmp.append(".tmp");
        read_program.pid = -1;
        read_program.read_fd = -1;
        progress_text = "0 bytes/Unknown";
        download_speed_text = "Unknown/s";
    }

    bool CurlDownloader::start() {
        remove(output_filepath_tmp.data.c_str());

        if(content_length == -1) {
            const char *args[] = { "curl",
                "-H", "Accept-Language: en-US,en;q=0.5", "-H", "Connection: keep-alive", "--compressed",
                "-H", "user-agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36",
                "--no-buffer", "-g", "-s", "-L", "-f", "-o", output_filepath_tmp.data.c_str(),
                "-D", "/dev/stdout",
                "--", url.c_str(), nullptr };

            if(exec_program_pipe(args, &read_program) != 0)
                return false;
        } else {
            const char *args[] = { "curl",
                "-H", "Accept-Language: en-US,en;q=0.5", "-H", "Connection: keep-alive", "--compressed",
                "-H", "user-agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36",
                "--no-buffer", "-g", "-s", "-L", "-f", "-o", output_filepath_tmp.data.c_str(),
                "--", url.c_str(), nullptr };

            if(exec_program_pipe(args, &read_program) != 0)
                return false;
        }

        if(content_length != -1)
            return true;

        // TODO: Remove this async task and make the fd non blocking instead
        header_reader = AsyncTask<bool>([this]{
            char tmp_buf[1024];
            while(true) {
                ssize_t bytes_available = read(read_program.read_fd, tmp_buf, sizeof(tmp_buf));
                if(bytes_available == -1) {
                    return false;
                } else if(bytes_available > 0 && content_length == -1) {
                    header.append(tmp_buf, bytes_available);
                    if(header.find("\r\n\r\n") != std::string::npos) {
                        std::string content_length_str = header_extract_value(header, "content-length");
                        if(!content_length_str.empty()) {
                            errno = 0;
                            char *endptr;
                            const int64_t content_length_tmp = strtoll(content_length_str.c_str(), &endptr, 10);
                            if(endptr != content_length_str.c_str() && errno == 0) {
                                std::lock_guard<std::mutex> lock(content_length_mutex);
                                content_length = content_length_tmp;
                            }
                        }
                    }
                }
            }
            return true;
        });

        return true;
    }

    bool CurlDownloader::stop(bool download_completed) {
        if(read_program.read_fd != -1)
            close(read_program.read_fd);
        if(read_program.pid != -1)
            kill(read_program.pid, SIGTERM);
        if(!download_completed)
            remove(output_filepath_tmp.data.c_str());
        //header_reader.cancel();
        finished = false;
        return true;
    }

    DownloadUpdateStatus CurlDownloader::update() {
        if(finished)
            return DownloadUpdateStatus::FINISHED;
        
        if(read_program.pid == -1)
            return DownloadUpdateStatus::ERROR;

        int status = 0;
        if(wait_program_non_blocking(read_program.pid, &status)) {
            read_program.pid = -1;
            if(status == 0 && rename_atomic(output_filepath_tmp.data.c_str(), output_filepath.c_str()) == 0) {
                finished = true;
                progress = 1.0f;
                return DownloadUpdateStatus::FINISHED;
            } else {
                return DownloadUpdateStatus::ERROR;
            }
        }

        if(header_reader.ready()) {
            if(!header_reader.get())
                return DownloadUpdateStatus::ERROR;
        }

        std::lock_guard<std::mutex> lock(content_length_mutex);
        size_t output_file_size = 0;
        file_get_size(output_filepath_tmp, &output_file_size);
        size_t downloaded_size = std::min((int64_t)output_file_size, content_length);

        if(content_length == -1) {
            progress_text = std::to_string(output_file_size / 1024) + "/Unknown";
        } else {
            size_t percentage = 0;
            if(output_file_size > 0)
                percentage = (double)downloaded_size / (double)content_length * 100.0;
            progress = (double)percentage / 100.0;
            progress_text = file_size_to_human_readable_string(downloaded_size) + "/" + file_size_to_human_readable_string(content_length) + " (" + std::to_string(percentage) + "%)";
        }

        // TODO: Take into consideration time overflow?
        size_t downloaded_diff = std::max(0lu, downloaded_size - downloaded_since_last_check);
        download_speed_text = file_size_to_human_readable_string(downloaded_diff) + "/s";
        downloaded_since_last_check = downloaded_size;

        return DownloadUpdateStatus::DOWNLOADING;
    }

    float CurlDownloader::get_progress() {
        return progress;
    }

    std::string CurlDownloader::get_progress_text() {
        return progress_text;
    }

    std::string CurlDownloader::get_download_speed_text() {
        return download_speed_text;
    }

    YoutubeDlDownloader::YoutubeDlDownloader(const std::string &url, const std::string &output_filepath, bool no_video) : Downloader(url, output_filepath), no_video(no_video) {
        // youtube-dl requires a file extension for the file
        if(this->output_filepath.find('.') == std::string::npos)
            this->output_filepath += ".mkv";

        read_program.pid = -1;
        read_program.read_fd = -1;
        progress_text = "0.0% of Unknown";
        download_speed_text = "Unknown/s";
    }

    bool YoutubeDlDownloader::start() {
        remove(output_filepath.c_str());

        std::vector<const char*> args = { "youtube-dl", "--no-warnings", "--no-continue", "--output", output_filepath.c_str(), "--newline" };
        if(no_video) {
            args.push_back("-f");
            args.push_back("bestaudio/best");
            args.push_back("-x");
        } else {
            args.push_back("-f");
            args.push_back("bestvideo+bestaudio/best");
        }
        args.insert(args.end(), { "--", url.c_str(), nullptr });

        if(exec_program_pipe(args.data(), &read_program) != 0)
            return false;

        read_program_file = fdopen(read_program.read_fd, "rb");
        if(!read_program_file) {
            if(read_program.pid != -1)
                wait_program(read_program.pid);
            return false;
        }

        // TODO: Remove this async task and make the fd non blocking instead
        youtube_dl_output_reader = AsyncTask<bool>([this]{
            char line[128];
            char progress_c[11];
            char content_size_c[21];
            char download_speed_c[21];

            while(true) {
                if(fgets(line, sizeof(line), read_program_file)) {
                    int len = strlen(line);
                    if(len > 0 && line[len - 1] == '\n') {
                        line[len - 1] = '\0';
                        --len;
                    }

                    if(sscanf(line, "[download] %10s of %20s at %20s", progress_c, content_size_c, download_speed_c) == 3) {
                        std::lock_guard<std::mutex> lock(progress_update_mutex);

                        if(strcmp(progress_c, "Unknown") != 0 && strcmp(content_size_c, "Unknown") != 0) {
                            std::string progress_str = progress_c;
                            progress_text = progress_str + " of " + content_size_c;
                            if(progress_str.back() == '%') {
                                errno = 0;
                                char *endptr;
                                const double progress_tmp = strtod(progress_str.c_str(), &endptr);
                                if(endptr != progress_str.c_str() && errno == 0)
                                    progress = progress_tmp / 100.0;
                            }
                        }

                        if(strcmp(download_speed_c, "Unknown") == 0)
                            download_speed_text = "Unknown/s";
                        else
                            download_speed_text = download_speed_c;
                    }
                } else {
                    return false;
                }
            }

            return true;
        });

        return true;
    }

    bool YoutubeDlDownloader::stop(bool) {
        if(read_program_file)
            fclose(read_program_file);
        if(read_program.pid != -1)
            kill(read_program.pid, SIGTERM);
        // TODO: Remove the temporary files created by youtube-dl (if !download_completed)
        //header_reader.cancel();
        finished = false;
        return true;
    }

    DownloadUpdateStatus YoutubeDlDownloader::update() {
        if(finished)
            return DownloadUpdateStatus::FINISHED;

        if(read_program.pid == -1)
            return DownloadUpdateStatus::ERROR;

        int status = 0;
        if(wait_program_non_blocking(read_program.pid, &status)) {
            read_program.pid = -1;
            if(status == 0) {
                finished = true;
                return DownloadUpdateStatus::FINISHED;
            } else {
                return DownloadUpdateStatus::ERROR;
            }
        }

        if(youtube_dl_output_reader.ready()) {
            if(!youtube_dl_output_reader.get())
                return DownloadUpdateStatus::ERROR;
        }

        return DownloadUpdateStatus::DOWNLOADING;
    }

    float YoutubeDlDownloader::get_progress() {
        std::lock_guard<std::mutex> lock(progress_update_mutex);
        return progress;
    }

    std::string YoutubeDlDownloader::get_progress_text() {
        std::lock_guard<std::mutex> lock(progress_update_mutex);
        return progress_text;
    }

    std::string YoutubeDlDownloader::get_download_speed_text() {
        std::lock_guard<std::mutex> lock(progress_update_mutex);
        return download_speed_text;
    }

    YoutubeDownloader::YoutubeDownloader(const MediaMetadata &video_metadata, const MediaMetadata &audio_metadata, const std::string &output_filepath) :
        Downloader("", output_filepath),
        video_metadata(video_metadata),
        audio_metadata(audio_metadata)
    {
        for(int i = 0; i < 2; ++i)
            downloaders[i] = nullptr;
    }

    bool YoutubeDownloader::start() {
        struct MediaProxyMetadata {
            std::unique_ptr<YoutubeMediaProxy> *media_proxy;
            MediaMetadata *media_metadata;
            std::unique_ptr<CurlDownloader> *downloader;
            std::string *output_filepath;
            bool tmp_audio_file;
        };

        const bool has_video = !video_metadata.url.empty();
        
        MediaProxyMetadata media_proxies[2] = {
            { &youtube_video_media_proxy, &video_metadata, &downloaders[0], &video_output_filepath, false },
            { &youtube_audio_media_proxy, &audio_metadata, &downloaders[1], &audio_output_filepath, has_video }
        };

        int num_proxied_media = 0;
        for(int i = 0; i < 2; ++i) {
            media_proxies[i].output_filepath->clear();
            if(media_proxies[i].media_metadata->url.empty() || youtube_url_is_live_stream(media_proxies[i].media_metadata->url))
                continue;

            *media_proxies[i].media_proxy = std::make_unique<YoutubeStaticMediaProxy>();
            if(!(*media_proxies[i].media_proxy)->start(media_proxies[i].media_metadata->url, media_proxies[i].media_metadata->content_length)) {
                fprintf(stderr, "Failed to load start youtube media proxy\n");
                return false;
            }

            std::string media_proxy_addr;
            if(!(*media_proxies[i].media_proxy)->get_address(media_proxy_addr)) {
                fprintf(stderr, "Failed to load start youtube media proxy\n");
                return false;
            }

            if(media_proxies[i].tmp_audio_file)
                *media_proxies[i].output_filepath = output_filepath + ".audio";
            else
                *media_proxies[i].output_filepath = output_filepath;

            remove(media_proxies[i].output_filepath->c_str());
            fprintf(stderr, "Downloading %s to %s\n", media_proxy_addr.c_str(), media_proxies[i].output_filepath->c_str());
            *media_proxies[i].downloader = std::make_unique<CurlDownloader>(media_proxy_addr, *media_proxies[i].output_filepath, media_proxies[i].media_metadata->content_length);
            *media_proxies[i].output_filepath = (*media_proxies[i].downloader)->get_output_filepath();
            if(!(*media_proxies[i].downloader)->start())
                return false;
            ++num_proxied_media;
        }

        if(num_proxied_media == 0) {
            return false;
        }

        downloader_task = AsyncTask<void>([this]() {
            // TODO: Poll instead of sleep
            while(!program_is_dead_in_current_thread()) {
                if(youtube_video_media_proxy)
                    youtube_video_media_proxy->update();

                if(youtube_audio_media_proxy)
                    youtube_audio_media_proxy->update();

                usleep(1000);
            }
        });

        return true;
    }

    static bool ffmpeg_merge_audio_and_video(const std::string &video_filepath, const std::string &audio_filepath, const std::string &output_filepath) {
        const char *args[] = { "ffmpeg", "-y", "-v", "quiet", "-i", video_filepath.c_str(), "-i", audio_filepath.c_str(), "-c", "copy", "--", output_filepath.c_str(), nullptr };
        return exec_program(args, nullptr, nullptr) == 0;
    }

    bool YoutubeDownloader::stop(bool download_completed) {
        for(int i = 0; i < 2; ++i) {
            if(downloaders[i]) {
                downloaders[i]->stop(download_completed);
                downloaders[i].reset();
            }
        }

        if(youtube_video_media_proxy)
            youtube_video_media_proxy->stop();

        if(youtube_audio_media_proxy)
            youtube_audio_media_proxy->stop();

        downloader_task.cancel();
        youtube_video_media_proxy.reset();
        youtube_audio_media_proxy.reset();

        if(!download_completed)
            return false;

        bool success = true;
        if(!video_metadata.url.empty() && !audio_metadata.url.empty()) {
            std::string tmp_file = video_output_filepath + ".tmp.mkv";
            if(ffmpeg_merge_audio_and_video(video_output_filepath, audio_output_filepath, tmp_file)) {
                success = rename_atomic(tmp_file.c_str(), output_filepath.c_str()) == 0;
                if(success) {
                    remove(audio_output_filepath.c_str());
                } else {
                    remove(tmp_file.c_str());
                }
            } else {
                success = false;
            }
        }
        return success;
    }

    DownloadUpdateStatus YoutubeDownloader::update() {
        int num_downloading = 0;
        int num_finished = 0;
        int num_failed = 0;
        int num_downloaders = 0;

        for(int i = 0; i < 2; ++i) {
            if(downloaders[i]) {
                ++num_downloaders;
                DownloadUpdateStatus update_status = downloaders[i]->update();
                switch(update_status) {
                    case DownloadUpdateStatus::DOWNLOADING:
                        ++num_downloading;
                        break;
                    case DownloadUpdateStatus::FINISHED:
                        ++num_finished;
                        break;
                    case DownloadUpdateStatus::ERROR:
                        ++num_failed;
                        break;
                }
            }
        }

        if(num_failed > 0)
            return DownloadUpdateStatus::ERROR;
        else if(num_finished == num_downloaders)
            return DownloadUpdateStatus::FINISHED;
        else
            return DownloadUpdateStatus::DOWNLOADING;
    }

    float YoutubeDownloader::get_progress() {
        return get_min_progress_downloader()->get_progress();
    }

    std::string YoutubeDownloader::get_progress_text() {
        return get_min_progress_downloader()->get_progress_text();
    }

    std::string YoutubeDownloader::get_download_speed_text() {
        return get_min_progress_downloader()->get_download_speed_text();
    }

    CurlDownloader* YoutubeDownloader::get_min_progress_downloader() {
        int min_index = -1;
        float progress = 999999.0f;
        for(int i = 0; i < 2; ++i) {
            if(downloaders[i]) {
                float item_progress = downloaders[i]->get_progress();
                if(item_progress < progress) {
                    min_index = i;
                    progress = item_progress;
                }
            }
        }
        //assert(min_index != -1);
        // Shouldn't happen
        if(min_index == -1)
            min_index = 0;
        return downloaders[min_index].get();
    }
}