#include    "Api.h"
#include    "Json.h"
#include    "curl/include/curl.h"

#include    <codecvt>
#include    <functional>
#include    <locale>

namespace {

    /**
     * Handle HTTP response
     * 
     * \param data HTTP response content
     * \param size Reponse data length
     * \param count Block count
     * \param userdata User data
     */
    size_t HandleResponse(void * data, size_t size, size_t count, void * userdata) {
        auto rsp = *(std::function<void (const std::string &)> *)userdata;
        rsp(std::string((char *)data, size * count));
        return size * count;
    }

    /**
     * Send HTTP GET request.
     * 
     * \param url URL to visit
     * \param authorization Token for HTTP header 'Authorization'
     * \param rsp Response handler
     */
    BOOL OpenUrl(const std::string & url, const std::string & authorization, const std::function<void (const std::string &)> & rsp) {
        CURL * ctx = curl_easy_init();
        curl_slist * header = NULL;

        if (ctx == nullptr) {
            printf("Fetch distro image failed. Due to curl_easy_init\n");
            return FALSE;
        }

        curl_easy_setopt(ctx, CURLOPT_POST, 0);
        curl_easy_setopt(ctx, CURLOPT_URL, url.c_str());
        curl_easy_setopt(ctx, CURLOPT_VERBOSE, 0);
        curl_easy_setopt(ctx, CURLOPT_WRITEFUNCTION, &HandleResponse);
        curl_easy_setopt(ctx, CURLOPT_WRITEDATA, &rsp);
        curl_easy_setopt(ctx, CURLOPT_HEADER, 0L);
        curl_easy_setopt(ctx, CURLOPT_SSL_VERIFYPEER, 0);
        curl_easy_setopt(ctx, CURLOPT_SSL_VERIFYHOST, 0);
        curl_easy_setopt(ctx, CURLOPT_NOSIGNAL, 1L);
        curl_easy_setopt(ctx, CURLOPT_FOLLOWLOCATION, 1L);

        if (!authorization.empty()) {
            header = curl_slist_append(header, authorization.c_str());
            curl_easy_setopt(ctx, CURLOPT_HTTPHEADER, header);
        }

        CURLcode ret = curl_easy_perform(ctx);
        curl_slist_free_all(header);
        if (ret != CURLE_OK) {
            printf("HTTP request failed. Due to curl_easy_perform ret : %d, %s\n", ret, curl_easy_strerror(ret));
            return FALSE;
        }

        return TRUE;
    }

    /**
     * Get user's input
     * 
     * \param message Input tips
     * \return User's input data
     */
    std::wstring GetUserInput(const std::wstring & message) {
        wprintf(L"%ls", message.c_str());

        wchar_t buf[512] = {0};
        std::wstring in;
        if (wscanf_s(L"%s", buf, 64) == 1) in = buf;
        
        /** Throw away any additional characters that did NOT fit in the buffer */
        wchar_t check;
        do {
            check = getwchar();
        } while ((check != L'\n') && check != WEOF);

        return in;
    }
}

Api::Api(const std::string & distro, const std::string & distro_repo) {
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    _distro = converter.from_bytes(distro);
    _repo = distro_repo;
    _dll = ::LoadLibraryEx(L"wslapi.dll", nullptr, LOAD_LIBRARY_SEARCH_SYSTEM32);
    if (_dll != nullptr) {
        _is_registered = (API_IS_REGISTERED)::GetProcAddress(_dll, "WslIsDistributionRegistered");
        _register = (API_REGISTER)::GetProcAddress(_dll, "WslRegisterDistribution");
        _unregister = (API_UNREGISTER)::GetProcAddress(_dll, "WslUnregisterDistribution");
        _config = (API_CONFIGURE)::GetProcAddress(_dll, "WslConfigureDistribution");
        _launch_interactive = (API_LAUNCH_INTERACTIVE)::GetProcAddress(_dll, "WslLaunchInteractive");
        _launch = (API_LAUNCH)::GetProcAddress(_dll, "WslLaunch");
    }
}

Api::~Api() {
    if (_dll != nullptr) FreeLibrary(_dll);
}

ULONG Api::GetUID(const std::wstring & name) const {
    HANDLE read_pipe;
    HANDLE write_pipe;

    SECURITY_ATTRIBUTES sa { sizeof(sa), nullptr, true };
    ULONG uid = (ULONG)-1;
    DWORD exit_code = 0;

    if (!::CreatePipe(&read_pipe, &write_pipe, &sa, 0)) {
        return uid;
    }

    std::wstring cmd = L"id -u " + name;
    HANDLE child;
    if (Launch(cmd.c_str(), true, ::GetStdHandle(STD_INPUT_HANDLE), write_pipe, GetStdHandle(STD_ERROR_HANDLE), &child)) {
        ::WaitForSingleObject(child, INFINITE);
        if (::GetExitCodeProcess(child, &exit_code) && (exit_code == 0)) {
            char buf[64] = {0};
            DWORD bytes = 0;

            if (::ReadFile(read_pipe, buf, 63, &bytes, nullptr)) {
                buf[bytes] = ANSI_NULL;
                try {
                    uid = std::stoul(buf, nullptr, 10);
                } catch (...) {}
            }
        }

        ::CloseHandle(child);
    }

    ::CloseHandle(read_pipe);
    ::CloseHandle(write_pipe);
    return uid;
}

BOOL Api::Install() const {
    printf("Installing, this may take a few minutes...\n");

    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::string repo = _repo;
    std::string token = "";

    /** Get Access token for Dockerhub */
    BOOL succ = OpenUrl("https://auth.docker.io/token?service=registry.docker.io&scope=repository:" + repo + ":pull", "", [&](const std::string & data) {
        Json::CharReaderBuilder builder;
        Json::Value rsp;

        if (!builder.newCharReader()->parse(data.data(), data.data() + data.size(), &rsp, nullptr) || !rsp.isObject() || !rsp["token"].isString()) {
            printf("Get access_token from auth.docker.io ... [Failed]\n", data.c_str());
            return;
        }

        token = "Authorization: Bearer " + rsp["token"].asString();
        printf("Get access_token from auth.docker.io ... [OK]\n");
    });
    if (!succ || token.empty()) return FALSE;
    
    /** Get all valid versions */
    std::vector<std::string> tags;
    succ = OpenUrl("https://registry.hub.docker.com/v2/" + repo + "/tags/list", token, [&](const std::string & data) {
        Json::CharReaderBuilder builder;
        Json::Value rsp;

        if (!builder.newCharReader()->parse(data.data(), data.data() + data.size(), &rsp, nullptr) || !rsp.isObject() || !rsp["tags"].isArray()) {
            printf("Get distro versions from registry.hub.docker.com ... [Failed]\n", data.c_str());
            return;
        }

        size_t count = rsp["tags"].size();
        if (count == 0) {
            printf("Bad distro. Can NOT find image from docker.io\n");
            return;
        }

        printf("Get distro versions from registry.hub.docker.com ... [OK]\n");
        printf("Valid versions for distro [%s]:\n", repo.c_str());
        for (int i = 0; i < count; ++i) {
            printf("\t%s\n", rsp["tags"][i].asCString());
            tags.push_back(rsp["tags"][i].asString());
        }
    });
    if (!succ || tags.empty()) return FALSE;

    /** Let user select which distro version to install */
    std::wstring input_tag = GetUserInput(L"Please input the version tag you want to install : ");
    std::string tag = converter.to_bytes(input_tag);

    /** Get selected version manifests */
    std::string rootfs;
    succ = OpenUrl("https://registry.hub.docker.com/v2/" + repo + "/manifests/" + tag, token, [&](const std::string & data) {
        Json::CharReaderBuilder builder;
        Json::Value rsp;

        if (!builder.newCharReader()->parse(data.data(), data.data() + data.size(), &rsp, nullptr) || !rsp.isObject() || !rsp["fsLayers"].isArray()) {
            printf("Get [%s:%s] manifest from registry.hub.docker.com ... [Failed]\n", repo.c_str(), tag.c_str(), data.c_str());
            return;
        }

        size_t count = rsp["fsLayers"].size();
        if (count == 0) {
            printf("Bad distro version : %s. Can NOT find image from docker.io\n", tag.c_str());
            return;
        }

        printf("Get [%s:%s] manifest from registry.hub.docker.com ... [OK]\n", repo.c_str(), tag.c_str());
        for (int i = (int)count - 1; i >= 0; --i) {
            std::string sum = rsp["fsLayers"][i]["blobSum"].asString();
            if (sum == "sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4") continue;
            rootfs = sum;
            break;
        }
    });
    if (!succ || rootfs.empty()) return FALSE;    

    /** Downloading file */
    {
        FILE * file = fopen("rootfs.tar.gz", "wb+");
        if (!file) {
            printf("Create file for downloading distro's image failed\n");
            return FALSE;
        }
        
        printf("Downloading rootfs for [%s:%s] ...\r", repo.c_str(), tag.c_str());
        size_t size = 0;
        bool ok = true;
        succ = OpenUrl("https://registry.hub.docker.com/v2/" + repo + "/blobs/" + rootfs, token, [&](const std::string & data) {
            if (data.size() > 0) {
                size_t writed = fwrite(data.data(), 1, data.size(), file);
                if (writed != data.size()) {
                    ok = false;
                } else {
                    size += writed;
                    printf("Downloading rootfs for [%s:%s] ... %.3lf MB\r", repo.c_str(), tag.c_str(), size * 1.0 / (1024 * 1024));
                }
            }
        });

        fflush(file);
        fclose(file);

        if (!ok) {
            printf("Downloading rootfs for [%s:%s] ... [WRITE ERROR]\n", repo.c_str(), tag.c_str());
            return FALSE;
        } else if (!succ) {
            printf("Downloading rootfs for [%s:%s] ... [FAILED]\n", repo.c_str(), tag.c_str());
            return FALSE;
        } else {
            printf("Downloading rootfs for [%s:%s] ... %.3lf MB [DONE]\n", repo.c_str(), tag.c_str(), size * 1.0 / (1024 * 1024));
        }
    }
    
    /** Register WSL */
    printf("Register distribution[%s:%s] ...\r", repo.c_str(), tag.c_str());
    HRESULT hr = _register(_distro.c_str(), L"rootfs.tar.gz");
    if (FAILED(hr)) {
        printf("Register distribution[%s:%s] ... [FAILED]\n", repo.c_str(), tag.c_str());
        return FALSE;
    } else {
        printf("Register distribution[%s:%s] ... [OK]\n", repo.c_str(), tag.c_str());
    }

    /** Delete /etc/resolv.conf to allow WSL to generate a version based on Windows network information */
    DWORD exit_code;
    hr = LaunchInteractive(L"/bin/rm /etc/resolv.conf", true, &exit_code);
    if (FAILED(hr)) {
        printf("Failed to install this distro : Can NOT rm /etc/resolv.conf\n");
        return FALSE;
    }

    /** Ask for default login user */
    CreateUser();
    
    /** Generate uninstall batch file */
    FILE * uninstall = fopen("uninstall.bat", "w+");
    std::string uninstall_cmd = converter.to_bytes(_distro) + ".exe uninstall\n";
    fwrite(uninstall_cmd.c_str(), 1, uninstall_cmd.size(), uninstall);
    fflush(uninstall);
    fclose(uninstall);

    /** Clean up */   
    remove("rootfs.tar.gz");
    return TRUE;
}

VOID Api::CreateUser() const {
    /** Confirm */
    while (true) {
        std::wstring create_user = GetUserInput(L"Create new default user instead of 'root'[y/n] : ");
        if (create_user == L"n") return;
        if (create_user == L"y") break;
    }

    /** Get user name */
    std::wstring name = GetUserInput(L"Create a default UNIX user. The username does not need to match your Windows username.\nFor more information visit: https://aka.ms/wslusers\nEnter new UNIX username: ");
    DWORD exit_code;

    /** Clean on failed */
    struct Cleaner {
        const Api * api;
        std::wstring name;
        bool cancel;

        ~Cleaner() {
            if (cancel) return;

            wprintf(L"Create default user failed, use 'root' as default\n");
            std::wstring cmd_deluser = L"userdel -f -r " + name;
            DWORD unused = 0;
            api->LaunchInteractive(cmd_deluser.c_str(), true, &unused);
        }

    } cleanup = { this, name, false };

    /** Create user account */
    std::wstring cmd_adduser = L"useradd -g root -G adm,wheel -m " + name;
    if (!LaunchInteractive(cmd_adduser.c_str(), true, &exit_code) || exit_code != 0) return;

    /** Set user's password */
    std::wstring cmd_pswd = L"passwd " + name;
    if (!LaunchInteractive(cmd_pswd.c_str(), true, &exit_code) || exit_code != 0) return;

    /** Set as default user */
    SetDefaultUser(name);
    cleanup.cancel = true;
}

VOID Api::Uninstall() const {
    DWORD exit_code;
    LaunchInteractive(L"/bin/rm -rf /usr", true, &exit_code);

    HRESULT hr = _unregister(_distro.c_str());
    if (FAILED(hr)) wprintf(L"Uninstall distribution[%ls] failed : %0x!\n", _distro.c_str(), hr);
    
    remove("rootfs");
}