#include "MCAuthCrypto.h"

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <bcrypt.h>
#pragma comment(lib, "bcrypt.lib")

#include <stdexcept>
#include <chrono>
#include <cassert>
#include <cstring>
#include <cstdio>

namespace MCAuth {

static void ThrowIfFailed(NTSTATUS status, const char* ctx) {
    if (!BCRYPT_SUCCESS(status)) {
        char buf[128];
        snprintf(buf, sizeof(buf), "BCrypt error 0x%08X in %s", (unsigned)status, ctx);
        throw std::runtime_error(buf);
    }
}

// Build SubjectPublicKeyInfo DER for an EC public key.
// Hardcoded for P-256 and P-384 (the only two we use).
std::vector<uint8_t> BuildSubjectPublicKeyInfoDER(
    int bitSize,
    const std::vector<uint8_t>& x,
    const std::vector<uint8_t>& y)
{
    // OIDs
    // ecPublicKey: 1.2.840.10045.2.1  -> 2a 86 48 ce 3d 02 01 (7 bytes)
    static const uint8_t OID_EC_PUBLIC_KEY[] = { 0x2a,0x86,0x48,0xce,0x3d,0x02,0x01 };
    // secp256r1 : 1.2.840.10045.3.1.7 -> 2a 86 48 ce 3d 03 01 07 (8 bytes)
    static const uint8_t OID_P256[] = { 0x2a,0x86,0x48,0xce,0x3d,0x03,0x01,0x07 };
    // secp384r1 : 1.3.132.0.34        -> 2b 81 04 00 22 (5 bytes)
    static const uint8_t OID_P384[] = { 0x2b,0x81,0x04,0x00,0x22 };

    size_t coordSize = (bitSize == 256) ? 32 : 48;
    const uint8_t* curveOid = (bitSize == 256) ? OID_P256 : OID_P384;
    size_t curveOidLen = (bitSize == 256) ? 8 : 5;

    // AlgorithmIdentifier SEQUENCE:
    //   06 07 <ecPublicKey OID>
    //   06 <len> <curve OID>
    size_t algIdInnerLen = 2 + sizeof(OID_EC_PUBLIC_KEY) + 2 + curveOidLen;

    // BIT STRING content: 00 04 X Y
    size_t bitStringContent = 1 + 1 + coordSize + coordSize; // 00 04 X Y
    size_t bitStringLen = bitStringContent;

    // Outer SEQUENCE content
    size_t outerContent = (2 + algIdInnerLen) + (2 + bitStringLen);

    // Single-byte DER lengths suffice for P-256 (89 bytes) and P-384 (118 bytes)

    std::vector<uint8_t> der;
    auto appendLen = [&](size_t len) {
        if (len < 128) {
            der.push_back((uint8_t)len);
        } else {
            der.push_back(0x82);
            der.push_back((uint8_t)(len >> 8));
            der.push_back((uint8_t)(len & 0xFF));
        }
    };

    // Outer SEQUENCE
    der.push_back(0x30);
    appendLen(outerContent);

    // AlgorithmIdentifier SEQUENCE
    der.push_back(0x30);
    appendLen(algIdInnerLen);
    der.push_back(0x06); der.push_back((uint8_t)sizeof(OID_EC_PUBLIC_KEY));
    der.insert(der.end(), OID_EC_PUBLIC_KEY, OID_EC_PUBLIC_KEY + sizeof(OID_EC_PUBLIC_KEY));
    der.push_back(0x06); der.push_back((uint8_t)curveOidLen);
    der.insert(der.end(), curveOid, curveOid + curveOidLen);

    // BIT STRING
    der.push_back(0x03);
    appendLen(bitStringLen + 1); // +1 for the "unused bits" byte
    der.push_back(0x00); // 0 unused bits
    der.push_back(0x04); // uncompressed point
    der.insert(der.end(), x.begin(), x.end());
    der.insert(der.end(), y.begin(), y.end());

    return der;
}

// Generate a key pair using BCrypt.
static ECKeyPair GenerateKeyPair(int bitSize) {
    LPCWSTR algId = (bitSize == 256) ? BCRYPT_ECDSA_P256_ALGORITHM
                                     : BCRYPT_ECDSA_P384_ALGORITHM;

    BCRYPT_ALG_HANDLE hAlg = nullptr;
    BCRYPT_KEY_HANDLE hKey = nullptr;

    try {
    ThrowIfFailed(BCryptOpenAlgorithmProvider(&hAlg, algId, nullptr, 0), "OpenAlgProvider");

    ThrowIfFailed(BCryptGenerateKeyPair(hAlg, &hKey, bitSize, 0), "GenerateKeyPair");
    ThrowIfFailed(BCryptFinalizeKeyPair(hKey, 0), "FinalizeKeyPair");

    // Export private blob
    ULONG privBlobLen = 0;
    ThrowIfFailed(BCryptExportKey(hKey, nullptr, BCRYPT_ECCPRIVATE_BLOB, nullptr, 0, &privBlobLen, 0), "ExportPriv(len)");
    std::vector<uint8_t> privBlob(privBlobLen);
    ThrowIfFailed(BCryptExportKey(hKey, nullptr, BCRYPT_ECCPRIVATE_BLOB, privBlob.data(), privBlobLen, &privBlobLen, 0), "ExportPriv");

    // Export public blob
    ULONG pubBlobLen = 0;
    ThrowIfFailed(BCryptExportKey(hKey, nullptr, BCRYPT_ECCPUBLIC_BLOB, nullptr, 0, &pubBlobLen, 0), "ExportPub(len)");
    std::vector<uint8_t> pubBlob(pubBlobLen);
    ThrowIfFailed(BCryptExportKey(hKey, nullptr, BCRYPT_ECCPUBLIC_BLOB, pubBlob.data(), pubBlobLen, &pubBlobLen, 0), "ExportPub");

    BCryptDestroyKey(hKey);
    BCryptCloseAlgorithmProvider(hAlg, 0);

    // BCRYPT_ECCKEY_BLOB header: 4-byte Magic + 4-byte cbKey
    // Public blob: [8-byte header] [X:cbKey bytes] [Y:cbKey bytes]
    const BCRYPT_ECCKEY_BLOB* hdr = reinterpret_cast<const BCRYPT_ECCKEY_BLOB*>(pubBlob.data());
    ULONG cbKey = hdr->cbKey;

    std::vector<uint8_t> x(pubBlob.begin() + 8,             pubBlob.begin() + 8 + cbKey);
    std::vector<uint8_t> y(pubBlob.begin() + 8 + cbKey,     pubBlob.begin() + 8 + cbKey * 2);

    ECKeyPair kp;
    kp.bitSize      = bitSize;
    kp.privateBlob  = std::move(privBlob);
    kp.publicBlob   = std::move(pubBlob);
    kp.x            = x;
    kp.y            = y;
    kp.publicKeyDER = BuildSubjectPublicKeyInfoDER(bitSize, x, y);
    return kp;

    } catch (...) {
        if (hKey) BCryptDestroyKey(hKey);
        if (hAlg) BCryptCloseAlgorithmProvider(hAlg, 0);
        throw;
    }
}

ECKeyPair GenerateP256KeyPair() { return GenerateKeyPair(256); }
ECKeyPair GenerateP384KeyPair() { return GenerateKeyPair(384); }

std::vector<uint8_t> SignSHA256P256(const ECKeyPair& kp,
                                     const uint8_t*   data,
                                     size_t           len)
{
    if (kp.bitSize != 256)
        throw std::runtime_error("SignSHA256P256: key is not P-256");

    BCRYPT_ALG_HANDLE hAlg = nullptr;
    BCRYPT_KEY_HANDLE hKey = nullptr;
    BCRYPT_ALG_HANDLE hHashAlg = nullptr;
    BCRYPT_HASH_HANDLE hHash = nullptr;

    try {
    ThrowIfFailed(BCryptOpenAlgorithmProvider(&hAlg, BCRYPT_ECDSA_P256_ALGORITHM, nullptr, 0), "Sign/OpenAlg");

    // Import private key
    ThrowIfFailed(BCryptImportKeyPair(hAlg, nullptr, BCRYPT_ECCPRIVATE_BLOB,
        &hKey,
        const_cast<uint8_t*>(kp.privateBlob.data()),
        (ULONG)kp.privateBlob.size(), 0), "Sign/ImportKey");

    // Hash the data with SHA-256
    ThrowIfFailed(BCryptOpenAlgorithmProvider(&hHashAlg, BCRYPT_SHA256_ALGORITHM, nullptr, 0), "Sign/OpenHashAlg");

    ThrowIfFailed(BCryptCreateHash(hHashAlg, &hHash, nullptr, 0, nullptr, 0, 0), "Sign/CreateHash");
    ThrowIfFailed(BCryptHashData(hHash, const_cast<uint8_t*>(data), (ULONG)len, 0), "Sign/HashData");

    uint8_t digest[32];
    ThrowIfFailed(BCryptFinishHash(hHash, digest, sizeof(digest), 0), "Sign/FinishHash");
    BCryptDestroyHash(hHash); hHash = nullptr;
    BCryptCloseAlgorithmProvider(hHashAlg, 0); hHashAlg = nullptr;

    // Sign the hash
    ULONG sigLen = 0;
    ThrowIfFailed(BCryptSignHash(hKey, nullptr, digest, sizeof(digest), nullptr, 0, &sigLen, 0), "Sign/GetSigLen");
    std::vector<uint8_t> sig(sigLen);
    ThrowIfFailed(BCryptSignHash(hKey, nullptr, digest, sizeof(digest), sig.data(), sigLen, &sigLen, 0), "Sign/SignHash");

    BCryptDestroyKey(hKey);
    BCryptCloseAlgorithmProvider(hAlg, 0);

    sig.resize(sigLen);
    return sig;

    } catch (...) {
        if (hHash) BCryptDestroyHash(hHash);
        if (hHashAlg) BCryptCloseAlgorithmProvider(hHashAlg, 0);
        if (hKey) BCryptDestroyKey(hKey);
        if (hAlg) BCryptCloseAlgorithmProvider(hAlg, 0);
        throw;
    }
}

static const char kB64Chars[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
static const char kB64UrlChars[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

static std::string Base64EncodeInternal(const uint8_t* data, size_t len, const char* alpha, bool pad) {
    std::string out;
    out.reserve(((len + 2) / 3) * 4);
    for (size_t i = 0; i < len; i += 3) {
        uint32_t v = (uint32_t)data[i] << 16;
        if (i + 1 < len) v |= (uint32_t)data[i + 1] << 8;
        if (i + 2 < len) v |= (uint32_t)data[i + 2];

        out += alpha[(v >> 18) & 0x3F];
        out += alpha[(v >> 12) & 0x3F];
        out += (i + 1 < len) ? alpha[(v >> 6) & 0x3F] : (pad ? '=' : '\0');
        out += (i + 2 < len) ? alpha[(v >> 0) & 0x3F] : (pad ? '=' : '\0');
    }
    // Remove trailing null chars (only happens for url encode without pad)
    while (!out.empty() && out.back() == '\0') out.pop_back();
    return out;
}

std::string Base64Encode(const uint8_t* data, size_t len) {
    return Base64EncodeInternal(data, len, kB64Chars, true);
}
std::string Base64Encode(const std::vector<uint8_t>& data) {
    return Base64Encode(data.data(), data.size());
}
std::string Base64UrlEncode(const uint8_t* data, size_t len) {
    return Base64EncodeInternal(data, len, kB64UrlChars, false);
}
std::string Base64UrlEncode(const std::vector<uint8_t>& data) {
    return Base64UrlEncode(data.data(), data.size());
}

std::string BuildProofKeyJson(const ECKeyPair& p256) {
    std::string xEnc = Base64UrlEncode(p256.x);
    std::string yEnc = Base64UrlEncode(p256.y);
    return "{\"kty\":\"EC\",\"alg\":\"ES256\",\"crv\":\"P-256\",\"use\":\"sig\","
           "\"x\":\"" + xEnc + "\",\"y\":\"" + yEnc + "\"}";
}

// XBL Signature header
// Payload to sign (all fields big-endian where integer):
//   INT32  policy_version = 1
//   BYTE   0x00
//   INT64  windows_timestamp
//   BYTE   0x00
//   BYTES  method  (e.g. "POST")
//   BYTE   0x00
//   BYTES  url_path_and_query
//   BYTE   0x00
//   BYTES  authorization_header  (empty string if none)
//   BYTE   0x00
//   BYTES  request_body
//   BYTE   0x00
//
// Signature header value (base64 of):
//   INT32  policy_version = 1
//   INT64  windows_timestamp
//   BYTES  64-byte P1363 ECDSA-SHA256 signature

static void AppendBE32(std::vector<uint8_t>& buf, uint32_t v) {
    buf.push_back((uint8_t)(v >> 24));
    buf.push_back((uint8_t)(v >> 16));
    buf.push_back((uint8_t)(v >>  8));
    buf.push_back((uint8_t)(v      ));
}
static void AppendBE64(std::vector<uint8_t>& buf, uint64_t v) {
    buf.push_back((uint8_t)(v >> 56));
    buf.push_back((uint8_t)(v >> 48));
    buf.push_back((uint8_t)(v >> 40));
    buf.push_back((uint8_t)(v >> 32));
    buf.push_back((uint8_t)(v >> 24));
    buf.push_back((uint8_t)(v >> 16));
    buf.push_back((uint8_t)(v >>  8));
    buf.push_back((uint8_t)(v      ));
}

std::string BuildXblSignatureHeader(const ECKeyPair&   p256,
                                     const std::string& method,
                                     const std::string& urlPath,
                                     const std::string& body,
                                     const std::string& authHdr)
{
    // Windows timestamp: (unix_seconds + 11644473600) * 10_000_000
    using namespace std::chrono;
    int64_t epochSec = duration_cast<seconds>(
        system_clock::now().time_since_epoch()).count();
    uint64_t winTs = (uint64_t)(epochSec + 11644473600LL) * 10000000ULL;

    // Build the content to sign
    std::vector<uint8_t> content;
    content.reserve(512);

    AppendBE32(content, 1);           // policy version
    content.push_back(0x00);
    AppendBE64(content, winTs);       // timestamp
    content.push_back(0x00);
    content.insert(content.end(), method.begin(), method.end());
    content.push_back(0x00);
    content.insert(content.end(), urlPath.begin(), urlPath.end());
    content.push_back(0x00);
    content.insert(content.end(), authHdr.begin(), authHdr.end());
    content.push_back(0x00);
    content.insert(content.end(), body.begin(), body.end());
    content.push_back(0x00);

    // Sign
    std::vector<uint8_t> sig = SignSHA256P256(p256, content.data(), content.size());

    // Build header blob
    std::vector<uint8_t> hdrBlob;
    hdrBlob.reserve(4 + 8 + sig.size());
    AppendBE32(hdrBlob, 1);
    AppendBE64(hdrBlob, winTs);
    hdrBlob.insert(hdrBlob.end(), sig.begin(), sig.end());

    return Base64Encode(hdrBlob);
}

std::string GenerateUUID() {
    uint8_t bytes[16];
    NTSTATUS st = BCryptGenRandom(nullptr, bytes, sizeof(bytes), BCRYPT_USE_SYSTEM_PREFERRED_RNG);
    if (!BCRYPT_SUCCESS(st)) {
        // fallback to rand (not cryptographic, but OK for device ID purposes)
        for (auto& b : bytes) b = (uint8_t)(rand() & 0xFF);
    }
    // Set version 4 and variant bits
    bytes[6] = (bytes[6] & 0x0F) | 0x40;
    bytes[8] = (bytes[8] & 0x3F) | 0x80;

    char buf[37];
    snprintf(buf, sizeof(buf),
        "%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x",
        bytes[0], bytes[1], bytes[2], bytes[3],
        bytes[4], bytes[5],
        bytes[6], bytes[7],
        bytes[8], bytes[9],
        bytes[10], bytes[11], bytes[12], bytes[13], bytes[14], bytes[15]);
    return buf;
}

std::vector<uint8_t> Base64Decode(const std::string& s) {
    static const int kInv[256] = {
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,62,-1,-1,-1,63,
        52,53,54,55,56,57,58,59,60,61,-1,-1,-1,-1,-1,-1,
        -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,10,11,12,13,14,
        15,16,17,18,19,20,21,22,23,24,25,-1,-1,-1,-1,-1,
        -1,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,
        41,42,43,44,45,46,47,48,49,50,51,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
        -1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1,
    };
    std::vector<uint8_t> out;
    int buf = 0, bits = 0;
    for (unsigned char c : s) {
        if (c == '=' || kInv[c] < 0) continue;
        buf  = (buf << 6) | kInv[c];
        bits += 6;
        if (bits >= 8) {
            bits -= 8;
            out.push_back((uint8_t)(buf >> bits));
            buf &= (1 << bits) - 1;
        }
    }
    return out;
}

std::string Base64DecodeStr(const std::string& encoded) {
    auto bytes = Base64Decode(encoded);
    return std::string(bytes.begin(), bytes.end());
}

std::vector<uint8_t> ComputeMD5(const void* data, size_t len) {
    std::vector<uint8_t> result(16, 0);
    BCRYPT_ALG_HANDLE hAlg = nullptr;
    BCRYPT_HASH_HANDLE hHash = nullptr;

    if (BCryptOpenAlgorithmProvider(&hAlg, BCRYPT_MD5_ALGORITHM, nullptr, 0) == 0) {
        if (BCryptCreateHash(hAlg, &hHash, nullptr, 0, nullptr, 0, 0) == 0) {
            if (BCryptHashData(hHash, (PUCHAR)data, (ULONG)len, 0) == 0) {
                BCryptFinishHash(hHash, result.data(), 16, 0);
            }
            BCryptDestroyHash(hHash);
        }
        BCryptCloseAlgorithmProvider(hAlg, 0);
    }
    return result;
}

} // namespace MCAuth