// clang-format off
#include <windows.h>
#include <shlobj.h>
#include <initguid.h>
#include <knownfolders.h>
// clang-format on

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <span>

#include "infra/infra.h"
#include "infra/sqlite_error.h"
#include "qmc2/qmc2.h"

using Infra::kgm_ekey_db_t;

constexpr std::array<uint8_t, 16> kMagicHeader{0x7C, 0xD5, 0x32, 0xEB, 0x86, 0x02, 0x7F, 0x4B,
                                               0xA8, 0xAF, 0xA6, 0x8E, 0x0F, 0xFF, 0x99, 0x14};

void DecryptKGG(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& kgg_file) {
    auto kgg_fname = kgg_file.filename().wstring();
    fwprintf(stderr, L"[INFO] reading %s...\n", kgg_fname.c_str());

    std::ifstream ifs_kgg(kgg_file, std::ios::binary);
    char header[0x100]{};
    ifs_kgg.read(header, sizeof(ifs_kgg));
    if (std::equal(kMagicHeader.cbegin(), kMagicHeader.cend(), header)) {
        fprintf(stderr, "[WARN] invalid kgg header, skip.\n");
        return;
    }
    uint32_t offset_to_audio = *reinterpret_cast<uint32_t*>(&header[0x10]);
    uint32_t encrypt_mode = *reinterpret_cast<uint32_t*>(&header[0x14]);
    if (encrypt_mode != 5) {
        fprintf(stderr, "[WARN] invalid enc_version (expect=0x05, got 0x%02x), skip.\n", encrypt_mode);
        return;
    }
    uint32_t audio_hash_len = *reinterpret_cast<uint32_t*>(&header[0x44]);
    if (audio_hash_len != 0x20) {
        fprintf(stderr, "[WARN] audio hash length invalid (expect=0x20, got 0x%02x), skip.\n", audio_hash_len);
        return;
    }
    std::string audio_hash(&header[0x48], &header[0x48 + audio_hash_len]);
    std::string ekey{};
    if (auto it = ekey_db.find(audio_hash); it != ekey_db.end()) {
        ekey = it->second;
    } else {
        fprintf(stderr, "[WARN] ekey not found, skip.\n");
        return;
    }

    auto qmc2 = QMC2::Create(ekey);
    if (!qmc2) {
        fprintf(stderr, "[WARN] create qmc2 failed, skip.\n");
        return;
    }

    ifs_kgg.seekg(offset_to_audio, std::ios::beg);
    std::vector<uint8_t> buffer(1024 * 1024);
    ifs_kgg.read(reinterpret_cast<char*>(buffer.data()), buffer.size());
    auto n = static_cast<size_t>(ifs_kgg.gcount());
    qmc2->Decrypt(std::span(buffer.begin(), n), 0);

    auto decrypted_path = kgg_file;
    auto magic = std::span(buffer.cbegin(), 4);
    if (std::equal(magic.begin(), magic.end(), "fLaC")) {
        decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.flac");
    } else if (std::equal(magic.begin(), magic.end(), "OggS")) {
        decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.ogg");
    } else {
        decrypted_path.replace_filename(decrypted_path.stem().wstring() + L"_kgg-dec.mp3");
    }
    if (exists(decrypted_path)) {
        fprintf(stderr, "[WARN] output file '%s' exists, skip.\n", decrypted_path.filename().string().c_str());
        return;
    }

    std::ofstream ofs_decrypted(decrypted_path, std::ios::binary);
    ofs_decrypted.write(reinterpret_cast<char*>(buffer.data()), n);
    size_t offset{n};
    while (!ifs_kgg.eof()) {
        ifs_kgg.read(reinterpret_cast<char*>(buffer.data()), buffer.size());
        n = static_cast<size_t>(ifs_kgg.gcount());
        qmc2->Decrypt(std::span(buffer.begin(), n), offset);
        ofs_decrypted.write(reinterpret_cast<char*>(buffer.data()), n);
        offset += n;
    }

    auto decrypted_fname = decrypted_path.filename().wstring();
    fwprintf(stderr, L"[INFO] OK! kgg --> %s\n", decrypted_fname.c_str());
}

void WalkFileOrDir(const kgm_ekey_db_t& ekey_db, const std::filesystem::path& input_path) {
    if (is_directory(input_path)) {
        for (auto const& dir_entry : std::filesystem::directory_iterator{input_path}) {
            DecryptKGG(ekey_db, dir_entry.path());
        }
    } else if (is_regular_file(input_path)) {
        DecryptKGG(ekey_db, input_path);
    }
}

int main() {
    SetConsoleOutputCP(CP_UTF8);
    setlocale(LC_ALL, ".UTF8");

    fputs("kgg-dec v0.1 by LSR\n", stderr);
    fputs("Usage: kgg-dec [kgg-dir=.]\n\n", stderr);
    PWSTR pAppDirPath{};
    SHGetKnownFolderPath(FOLDERID_RoamingAppData, 0, nullptr, &pAppDirPath);
    std::filesystem::path kgmPath{pAppDirPath};
    CoTaskMemFree(pAppDirPath);
    kgmPath = kgmPath / L"Kugou8" / L"KGMusicV3.db";

    int error{-1};
    Infra::KugouDb db{L"infra.dll", kgmPath};
    if (!db.IsOpen()) {
        fprintf(stderr, "[ERR ] db init error: is infra.dll ok?\n");
        return 1;
    }
    auto ekey_db = db.dump_ekey(error);
    if (error != 0) {
        fprintf(stderr, "[ERR ] dump ekey failed %d (%s)", error, sqlite_get_error(error));
        return 1;
    }
    db.Close();

#ifdef _DEBUG
    printf("ekey_db:\n");
    for (auto& [a, b] : ekey_db) {
        printf("%s --> %s\n", a.c_str(), b.c_str());
    }
#endif

    int argc;
    wchar_t** argv = CommandLineToArgvW(GetCommandLineW(), &argc);

    std::filesystem::path input_dir{(argc < 2) ? L"." : argv[1]};
    input_dir = absolute(input_dir);
    WalkFileOrDir(ekey_db, input_dir);

    return 0;
}