#include "../../plugins/Youtube.hpp"
#include "../../include/Storage.hpp"
#include <json/reader.h>
#include <json/writer.h>
#include <string.h>
#include <unordered_set>

namespace QuickMedia {
    static void iterate_suggestion_result(const Json::Value &value, std::vector<std::string> &result_items, int &iterate_count) {
        ++iterate_count;
        if(value.isArray()) {
            for(const Json::Value &child : value) {
                iterate_suggestion_result(child, result_items, iterate_count);
            }
        } else if(value.isString() && iterate_count > 2) {
            result_items.push_back(value.asString());
        }
    }

    static std::shared_ptr<BodyItem> parse_content_video_renderer(const Json::Value &content_item_json, std::unordered_set<std::string> &added_videos) {
        if(!content_item_json.isObject())
            return nullptr;
        
        const Json::Value &video_renderer_json = content_item_json["videoRenderer"];
        if(!video_renderer_json.isObject())
            return nullptr;
        
        const Json::Value &video_id_json = video_renderer_json["videoId"];
        if(!video_id_json.isString())
            return nullptr;

        std::string video_id_str = video_id_json.asString();
        if(added_videos.find(video_id_str) != added_videos.end())
            return nullptr;

        std::string thumbnail_url = "https://img.youtube.com/vi/" + video_id_str + "/hqdefault.jpg";

        const char *date = nullptr;
        const Json::Value &published_time_text_json = video_renderer_json["publishedTimeText"];
        if(published_time_text_json.isObject()) {
            const Json::Value &text_json = published_time_text_json["simpleText"];
            if(text_json.isString())
                date = text_json.asCString();
        }

        const char *length = nullptr;
        const Json::Value &length_text_json = video_renderer_json["lengthText"];
        if(length_text_json.isObject()) {
            const Json::Value &text_json = length_text_json["simpleText"];
            if(text_json.isString())
                length = text_json.asCString();
        }

        const char *title = nullptr;
        const Json::Value &title_json = video_renderer_json["title"];
        if(title_json.isObject()) {
            const Json::Value &runs_json = title_json["runs"];
            if(runs_json.isArray() && !runs_json.empty()) {
                const Json::Value &first_runs_json = runs_json[0];
                if(first_runs_json.isObject()) {
                    const Json::Value &text_json = first_runs_json["text"];
                    if(text_json.isString())
                        title = text_json.asCString();
                }
            }
        }

        if(!title)
            return nullptr;
        
        auto body_item = BodyItem::create(title);
        /* TODO: Make date a different color */
        std::string date_str;
        if(date)
            date_str += date;
        if(length) {
            if(!date_str.empty())
                date_str += '\n';
            date_str += length;
        }
        body_item->set_description(std::move(date_str));
        body_item->url = "https://www.youtube.com/watch?v=" + video_id_str;
        body_item->thumbnail_url = std::move(thumbnail_url);
        added_videos.insert(video_id_str);
        return body_item;
    }

    Youtube::Youtube() : Plugin("youtube") {
        
    }

    PluginResult Youtube::get_front_page(BodyItems &result_items) {
        bool disabled = true;
        if(disabled)
            return PluginResult::OK;

        std::string url = "https://youtube.com/";

        std::vector<CommandArg> additional_args = {
            { "-H", "x-spf-referer: " + url },
            { "-H", "x-youtube-client-name: 1" },
            { "-H", "x-youtube-client-version: 2.20200626.03.00" },
            { "-H", "referer: " + url }
        };

        std::vector<CommandArg> cookies = get_cookies();
        additional_args.insert(additional_args.end(), cookies.begin(), cookies.end());

        std::string website_data;
        if(download_to_string(url + "?pbj=1", website_data, additional_args, use_tor, true) != DownloadResult::OK)
            return PluginResult::NET_ERR;

        Json::Value json_root;
        Json::CharReaderBuilder json_builder;
        std::unique_ptr<Json::CharReader> json_reader(json_builder.newCharReader());
        std::string json_errors;
        if(!json_reader->parse(&website_data[0], &website_data[website_data.size()], &json_root, &json_errors)) {
            fprintf(stderr, "Youtube get front page error: %s\n", json_errors.c_str());
            return PluginResult::ERR;
        }

        if(!json_root.isArray())
            return PluginResult::ERR;

        std::unordered_set<std::string> added_videos;

        for(const Json::Value &json_item : json_root) {
            if(!json_item.isObject())
                continue;

            const Json::Value &response_json = json_item["response"];
            if(!response_json.isObject())
                continue;

            const Json::Value &contents_json = response_json["contents"];
            if(!contents_json.isObject())
                continue;

            const Json::Value &tcbrr_json = contents_json["twoColumnBrowseResultsRenderer"];
            if(!tcbrr_json.isObject())
                continue;

            const Json::Value &tabs_json = tcbrr_json["tabs"];
            if(!tabs_json.isArray())
                continue;

            for(const Json::Value &tab_item_json : tabs_json) {
                if(!tab_item_json.isObject())
                    continue;
                
                const Json::Value &tab_renderer_json = tab_item_json["tabRenderer"];
                if(!tab_renderer_json.isObject())
                    continue;
                
                const Json::Value &content_json = tab_renderer_json["content"];
                if(!content_json.isObject())
                    continue;

                const Json::Value &rich_grid_renderer = content_json["richGridRenderer"];
                if(!rich_grid_renderer.isObject())
                    continue;

                const Json::Value &contents2_json = rich_grid_renderer["contents"];
                if(!contents2_json.isArray())
                    continue;

                for(const Json::Value &contents_item : contents2_json) {
                    const Json::Value &rich_item_renderer_json = contents_item["richItemRenderer"];
                    if(!rich_item_renderer_json.isObject())
                        continue;
                    
                    const Json::Value &rich_item_contents = rich_item_renderer_json["content"];
                    std::shared_ptr<BodyItem> body_item = parse_content_video_renderer(rich_item_contents, added_videos);
                    if(body_item)
                        result_items.push_back(std::move(body_item));
                }
            }
        }

        return PluginResult::OK;
    }

    std::string Youtube::autocomplete_search(const std::string &query) {
        // Return the last result if the query is a substring of the autocomplete result
        if(last_autocomplete_result.size() >= query.size() && memcmp(query.data(), last_autocomplete_result.data(), query.size()) == 0)
            return last_autocomplete_result;

        std::string url = "https://clients1.google.com/complete/search?client=youtube&hl=en&gs_rn=64&gs_ri=youtube&ds=yt&cp=7&gs_id=x&q=";
        url += url_param_encode(query);

        std::string server_response;
        if(download_to_string(url, server_response, {}, use_tor, true) != DownloadResult::OK)
            return query;

        size_t json_start = server_response.find_first_of('(');
        if(json_start == std::string::npos)
            return query;
        ++json_start;

        size_t json_end = server_response.find_last_of(')');
        if(json_end == std::string::npos)
            return query;

        if(json_end == 0 || json_start >= json_end)
            return query;

        Json::Value json_root;
        Json::CharReaderBuilder json_builder;
        std::unique_ptr<Json::CharReader> json_reader(json_builder.newCharReader());
        std::string json_errors;
        if(!json_reader->parse(&server_response[json_start], &server_response[json_end], &json_root, &json_errors)) {
            fprintf(stderr, "Youtube autocomplete search json error: %s\n", json_errors.c_str());
            return query;
        }

        int iterate_count = 0;
        std::vector<std::string> result_items;
        iterate_suggestion_result(json_root, result_items, iterate_count);
        if(result_items.empty())
            return query;
        
        last_autocomplete_result = result_items[0];
        return result_items[0];
    }

    // Returns empty string if continuation token can't be found
    static std::string item_section_renderer_get_continuation_token(const Json::Value &item_section_renderer_json) {
        const Json::Value &continuations_json = item_section_renderer_json["continuations"];
        if(!continuations_json.isArray() || continuations_json.empty())
            return "";
        
        const Json::Value &first_continuation_json = continuations_json[0];
        if(!first_continuation_json.isObject())
            return "";
        
        const Json::Value &next_continuation_data_json = first_continuation_json["nextContinuationData"];
        if(!next_continuation_data_json.isObject())
            return "";
        
        const Json::Value &continuation_json = next_continuation_data_json["continuation"];
        if(!continuation_json.isString())
            return "";
        
        return continuation_json.asString();
    }

    static void parse_item_section_renderer(const Json::Value &item_section_renderer_json, std::string &continuation_token, std::unordered_set<std::string> &added_videos, BodyItems &result_items) {
        if(continuation_token.empty())
            continuation_token = item_section_renderer_get_continuation_token(item_section_renderer_json);
        
        const Json::Value &item_contents_json = item_section_renderer_json["contents"];
        if(!item_contents_json.isArray())
            return;

        for(const Json::Value &content_item_json : item_contents_json) {
            if(!content_item_json.isObject())
                continue;

            const Json::Value &shelf_renderer_json = content_item_json["shelfRenderer"];
            if(!shelf_renderer_json.isObject())
                continue;

            const Json::Value &item_content_json = shelf_renderer_json["content"];
            if(!item_content_json.isObject())
                continue;

            const Json::Value &vertical_list_renderer_json = item_content_json["verticalListRenderer"];
            if(!vertical_list_renderer_json.isObject())
                continue;

            const Json::Value &items_json = vertical_list_renderer_json["items"];
            if(!items_json.isArray())
                continue;
            
            for(const Json::Value &item_json : items_json) {
                std::shared_ptr<BodyItem> body_item = parse_content_video_renderer(item_json, added_videos);
                if(body_item)
                    result_items.push_back(std::move(body_item));
            }
        }
    
        for(const Json::Value &content_item_json : item_contents_json) {
            std::shared_ptr<BodyItem> body_item = parse_content_video_renderer(content_item_json, added_videos);
            if(body_item)
                result_items.push_back(std::move(body_item));
        }
    }

    SuggestionResult Youtube::update_search_suggestions(const std::string &text, BodyItems &result_items) {
        std::string url = "https://youtube.com/results?search_query=";
        url += url_param_encode(text);

        std::vector<CommandArg> additional_args = {
            { "-H", "x-spf-referer: " + url },
            { "-H", "x-youtube-client-name: 1" },
            { "-H", "x-youtube-client-version: 2.20200626.03.00" },
            { "-H", "referer: " + url }
        };

        std::vector<CommandArg> cookies = get_cookies();
        additional_args.insert(additional_args.end(), cookies.begin(), cookies.end());

        std::string website_data;
        if(download_to_string(url + "&pbj=1", website_data, additional_args, use_tor, true) != DownloadResult::OK)
            return SuggestionResult::NET_ERR;

        Json::Value json_root;
        Json::CharReaderBuilder json_builder;
        std::unique_ptr<Json::CharReader> json_reader(json_builder.newCharReader());
        std::string json_errors;
        if(!json_reader->parse(&website_data[0], &website_data[website_data.size()], &json_root, &json_errors)) {
            fprintf(stderr, "Youtube search json error: %s\n", json_errors.c_str());
            return SuggestionResult::ERR;
        }

        if(!json_root.isArray())
            return SuggestionResult::ERR;

        std::string continuation_token;
        std::unordered_set<std::string> added_videos; /* The input contains duplicates, filter them out! */

        for(const Json::Value &json_item : json_root) {
            if(!json_item.isObject())
                continue;

            const Json::Value &response_json = json_item["response"];
            if(!response_json.isObject())
                continue;

            const Json::Value &contents_json = response_json["contents"];
            if(!contents_json.isObject())
                continue;

            const Json::Value &tcsrr_json = contents_json["twoColumnSearchResultsRenderer"];
            if(!tcsrr_json.isObject())
                continue;

            const Json::Value &primary_contents_json = tcsrr_json["primaryContents"];
            if(!primary_contents_json.isObject())
                continue;

            const Json::Value &section_list_renderer_json = primary_contents_json["sectionListRenderer"];
            if(!section_list_renderer_json.isObject())
                continue;
            
            const Json::Value &contents2_json = section_list_renderer_json["contents"];
            if(!contents2_json.isArray())
                continue;
            
            for(const Json::Value &item_json : contents2_json) {
                if(!item_json.isObject())
                    continue;
                
                const Json::Value &item_section_renderer_json = item_json["itemSectionRenderer"];
                if(!item_section_renderer_json.isObject())
                    continue;

                parse_item_section_renderer(item_section_renderer_json, continuation_token, added_videos, result_items);
            }
        }

        // The continuation data can also contain continuation, but we ignore that for now. Only get the first continuation data
        if(!continuation_token.empty())
            search_suggestions_get_continuation(url, continuation_token, result_items);

        return SuggestionResult::OK;
    }

    void Youtube::search_suggestions_get_continuation(const std::string &url, const std::string &continuation_token, BodyItems &result_items) {
        std::string next_url = url + "&pbj=1&ctoken=" + continuation_token;

        std::vector<CommandArg> additional_args = {
            { "-H", "x-spf-referer: " + url },
            { "-H", "x-youtube-client-name: 1" },
            { "-H", "x-spf-previous: " + url },
            { "-H", "x-youtube-client-version: 2.20200626.03.00" },
            { "-H", "referer: " + url }
        };

        std::vector<CommandArg> cookies = get_cookies();
        additional_args.insert(additional_args.end(), cookies.begin(), cookies.end());

        std::string website_data;
        if(download_to_string(next_url, website_data, additional_args, use_tor, true) != DownloadResult::OK)
            return;

        Json::Value json_root;
        Json::CharReaderBuilder json_builder;
        std::unique_ptr<Json::CharReader> json_reader(json_builder.newCharReader());
        std::string json_errors;
        if(!json_reader->parse(&website_data[0], &website_data[website_data.size()], &json_root, &json_errors)) {
            fprintf(stderr, "Youtube search continuation json error: %s\n", json_errors.c_str());
            return;
        }

        if(!json_root.isArray())
            return;

        std::string next_continuation_token;
        std::unordered_set<std::string> added_videos;

        for(const Json::Value &json_item : json_root) {
            if(!json_item.isObject())
                continue;

            const Json::Value &response_json = json_item["response"];
            if(!response_json.isObject())
                continue;

            const Json::Value &continuation_contents_json = response_json["continuationContents"];
            if(!continuation_contents_json.isObject())
                continue;

            const Json::Value &item_section_continuation_json = continuation_contents_json["itemSectionContinuation"];
            if(!item_section_continuation_json.isObject())
                continue;

            // Note: item_section_continuation json object is compatible with item_section_renderer json object
            parse_item_section_renderer(item_section_continuation_json, next_continuation_token, added_videos, result_items);
        }
    }

    std::vector<CommandArg> Youtube::get_cookies() const {
        if(use_tor)
            return {};

        Path cookies_filepath;
        if(get_cookies_filepath(cookies_filepath, name) != 0) {
            fprintf(stderr, "Warning: Failed to create youtube cookies file\n");
            return {};
        }

        return {
            CommandArg{ "-b", cookies_filepath.data },
            CommandArg{ "-c", cookies_filepath.data }
        };
    }

    static std::string remove_index_from_playlist_url(const std::string &url) {
        std::string result = url;
        size_t index = result.rfind("&index=");
        if(index == std::string::npos)
            return result;
        return result.substr(0, index);
    }

    static std::shared_ptr<BodyItem> parse_compact_video_renderer_json(const Json::Value &item_json, std::unordered_set<std::string> &added_videos) {
        const Json::Value &compact_video_renderer_json = item_json["compactVideoRenderer"];
        if(!compact_video_renderer_json.isObject())
            return nullptr;
        
        const Json::Value &video_id_json = compact_video_renderer_json["videoId"];
        if(!video_id_json.isString())
            return nullptr;

        std::string video_id_str = video_id_json.asString();
        if(added_videos.find(video_id_str) != added_videos.end())
            return nullptr;

        std::string thumbnail_url = "https://img.youtube.com/vi/" + video_id_str + "/hqdefault.jpg";

        const char *date = nullptr;
        const Json::Value &published_time_text_json = compact_video_renderer_json["publishedTimeText"];
        if(published_time_text_json.isObject()) {
            const Json::Value &text_json = published_time_text_json["simpleText"];
            if(text_json.isString())
                date = text_json.asCString();
        }

        const char *length = nullptr;
        const Json::Value &length_text_json = compact_video_renderer_json["lengthText"];
        if(length_text_json.isObject()) {
            const Json::Value &text_json = length_text_json["simpleText"];
            if(text_json.isString())
                length = text_json.asCString();
        }

        const char *title = nullptr;
        const Json::Value &title_json = compact_video_renderer_json["title"];
        if(title_json.isObject()) {
            const Json::Value &simple_text_json = title_json["simpleText"];
            if(simple_text_json.isString()) {
                title = simple_text_json.asCString();
            }
        }

        if(!title)
            return nullptr;
        
        auto body_item = BodyItem::create(title);
        /* TODO: Make date a different color */
        std::string date_str;
        if(date)
            date_str += date;
        if(length) {
            if(!date_str.empty())
                date_str += '\n';
            date_str += length;
        }
        body_item->set_description(std::move(date_str));
        body_item->url = "https://www.youtube.com/watch?v=" + video_id_str;
        body_item->thumbnail_url = std::move(thumbnail_url);
        added_videos.insert(video_id_str);
        return body_item;
    }

    // TODO: Make this faster by using string search instead of parsing html.
    // TODO: If the result is a play
    BodyItems Youtube::get_related_media(const std::string &url) {
        BodyItems result_items;

        std::string modified_url = remove_index_from_playlist_url(url);
        std::vector<CommandArg> additional_args = {
            { "-H", "x-spf-referer: " + url },
            { "-H", "x-youtube-client-name: 1" },
            { "-H", "x-spf-previous: " + url },
            { "-H", "x-youtube-client-version: 2.20200626.03.00" },
            { "-H", "referer: " + url }
        };

        std::vector<CommandArg> cookies = get_cookies();
        additional_args.insert(additional_args.end(), cookies.begin(), cookies.end());

        std::string website_data;
        if(download_to_string(modified_url + "&pbj=1", website_data, additional_args, use_tor, true) != DownloadResult::OK)
            return result_items;

        Json::Value json_root;
        Json::CharReaderBuilder json_builder;
        std::unique_ptr<Json::CharReader> json_reader(json_builder.newCharReader());
        std::string json_errors;
        if(!json_reader->parse(&website_data[0], &website_data[website_data.size()], &json_root, &json_errors)) {
            fprintf(stderr, "Youtube related media error: %s\n", json_errors.c_str());
            return result_items;
        }

        if(!json_root.isArray())
            return result_items;

        std::unordered_set<std::string> added_videos;

        for(const Json::Value &json_item : json_root) {
            if(!json_item.isObject())
                continue;

            const Json::Value &response_json = json_item["response"];
            if(!response_json.isObject())
                continue;

            const Json::Value &contents_json = response_json["contents"];
            if(!contents_json.isObject())
                return result_items;

            const Json::Value &tcwnr_json = contents_json["twoColumnWatchNextResults"];
            if(!tcwnr_json.isObject())
                return result_items;

            const Json::Value &secondary_results_json = tcwnr_json["secondaryResults"];
            if(!secondary_results_json.isObject())
                return result_items;

            const Json::Value &secondary_results2_json = secondary_results_json["secondaryResults"];
            if(!secondary_results2_json.isObject())
                return result_items;
            
            const Json::Value &results_json = secondary_results2_json["results"];
            if(!results_json.isArray())
                return result_items;

            for(const Json::Value &item_json : results_json) {
                if(!item_json.isObject())
                    continue;

                auto body_item = parse_compact_video_renderer_json(item_json, added_videos);
                if(body_item)
                    result_items.push_back(std::move(body_item));
                
                const Json::Value &compact_autoplay_renderer_json = item_json["compactAutoplayRenderer"];
                if(!compact_autoplay_renderer_json.isObject())
                    continue;
                
                const Json::Value &item_contents_json = compact_autoplay_renderer_json["contents"];
                if(!item_contents_json.isArray())
                    continue;
                
                for(const Json::Value &content_item_json : item_contents_json) {
                    if(!content_item_json.isObject())
                        continue;
                    
                    auto body_item = parse_compact_video_renderer_json(content_item_json, added_videos);
                    if(body_item)
                        result_items.push_back(std::move(body_item));
                }
            }
        }

        return result_items;
    }
}