#include "TokenStore.h"
#include <shlobj.h>
#include <cstdio>
#include <cstring>
#include <vector>
#include <fstream>

#pragma comment(lib, "crypt32.lib")
#pragma comment(lib, "shell32.lib")

// ─── Singleton ────────────────────────────────────────

TokenStore* TokenStore::Get()
{
    static TokenStore instance;
    return &instance;
}

// ─── Storage Path ─────────────────────────────────────

std::string TokenStore::GetStoragePath()
{
    char appData[MAX_PATH];
    if (SUCCEEDED(SHGetFolderPathA(nullptr, CSIDL_APPDATA, nullptr, 0, appData)))
    {
        std::string dir = std::string(appData) + "\\skeet_loader";
        CreateDirectoryA(dir.c_str(), nullptr);
        return dir + "\\.session";
    }
    // Fallback: current directory
    return ".session";
}

// ─── DPAPI Encrypt ────────────────────────────────────

std::vector<BYTE> TokenStore::Encrypt(const std::vector<BYTE>& plaintext)
{
    DATA_BLOB in, out;
    in.pbData = const_cast<BYTE*>(plaintext.data());
    in.cbData = static_cast<DWORD>(plaintext.size());

    if (!CryptProtectData(&in, L"skeet_loader_session", nullptr, nullptr, nullptr,
                          CRYPTPROTECT_UI_FORBIDDEN | CRYPTPROTECT_LOCAL_MACHINE, &out))
    {
        printf("[TokenStore] CryptProtectData failed: 0x%lx\n", GetLastError());
        return {};
    }

    std::vector<BYTE> result(out.pbData, out.pbData + out.cbData);
    LocalFree(out.pbData);
    return result;
}

// ─── DPAPI Decrypt ────────────────────────────────────

std::vector<BYTE> TokenStore::Decrypt(const std::vector<BYTE>& ciphertext)
{
    DATA_BLOB in, out;
    in.pbData = const_cast<BYTE*>(ciphertext.data());
    in.cbData = static_cast<DWORD>(ciphertext.size());

    if (!CryptUnprotectData(&in, nullptr, nullptr, nullptr, nullptr,
                            CRYPTPROTECT_UI_FORBIDDEN, &out))
    {
        printf("[TokenStore] CryptUnprotectData failed: 0x%lx\n", GetLastError());
        return {};
    }

    std::vector<BYTE> result(out.pbData, out.pbData + out.cbData);
    LocalFree(out.pbData);
    return result;
}

// ─── Public API ───────────────────────────────────────

bool TokenStore::Save(const std::string& refreshToken, const std::string& username)
{
    // Build simple JSON
    std::string json = "{\"refresh_token\":\"" + refreshToken + "\",\"username\":\"" + username + "\"}";

    std::vector<BYTE> plaintext(json.begin(), json.end());
    auto encrypted = Encrypt(plaintext);

    if (encrypted.empty())
    {
        printf("[TokenStore] Save: encryption failed\n");
        return false;
    }

    std::string path = GetStoragePath();
    std::ofstream file(path, std::ios::binary);
    if (!file)
    {
        printf("[TokenStore] Save: cannot open %s\n", path.c_str());
        return false;
    }

    file.write(reinterpret_cast<const char*>(encrypted.data()), encrypted.size());
    file.close();

    printf("[TokenStore] Saved session to %s (%zu bytes encrypted)\n", path.c_str(), encrypted.size());
    return true;
}

TokenStore::SessionData TokenStore::Load()
{
    SessionData data;

    std::string path = GetStoragePath();
    std::ifstream file(path, std::ios::binary | std::ios::ate);
    if (!file)
    {
        printf("[TokenStore] Load: no saved session at %s\n", path.c_str());
        return data;
    }

    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);

    std::vector<BYTE> encrypted(size);
    file.read(reinterpret_cast<char*>(encrypted.data()), size);
    file.close();

    auto decrypted = Decrypt(encrypted);
    if (decrypted.empty())
    {
        printf("[TokenStore] Load: decryption failed\n");
        // Clean up corrupted file
        Clear();
        return data;
    }

    std::string json(decrypted.begin(), decrypted.end());
    printf("[TokenStore] Load: decrypted %zu bytes\n", decrypted.size());

    // Simple JSON parsing (no dependency needed)
    auto extract = [&json](const std::string& key) -> std::string {
        std::string search = "\"" + key + "\":\"";
        size_t start = json.find(search);
        if (start == std::string::npos) return {};
        start += search.length();
        size_t end = json.find("\"", start);
        if (end == std::string::npos) return {};
        return json.substr(start, end - start);
    };

    data.refreshToken = extract("refresh_token");
    data.username = extract("username");
    data.valid = !data.refreshToken.empty();

    printf("[TokenStore] Load: username=%s valid=%d\n", data.username.c_str(), data.valid);
    return data;
}

void TokenStore::Clear()
{
    std::string path = GetStoragePath();
    DeleteFileA(path.c_str());
    printf("[TokenStore] Cleared session: %s\n", path.c_str());
}
