#include "crypto_api.h"

#include <lua.hpp>
#include <sstream>
#include <iomanip>
#include <cstring>

#ifdef _WIN32
#include <windows.h>
#include <bcrypt.h>
#pragma comment(lib, "bcrypt.lib")
#elif defined(MOSIS_HAS_OPENSSL)
#include <openssl/evp.h>
#include <openssl/hmac.h>
#endif

namespace mosis {

//=============================================================================
// SECURE RANDOM
//=============================================================================

SecureRandom::SecureRandom()
    : m_gen(m_rd()) {
}

std::string SecureRandom::GetBytes(size_t count) {
    std::lock_guard<std::mutex> lock(m_mutex);

    std::string result(count, '\0');
    for (size_t i = 0; i < count; i++) {
        result[i] = static_cast<char>(m_gen() & 0xFF);
    }
    return result;
}

int64_t SecureRandom::GetInt(int64_t min, int64_t max) {
    std::lock_guard<std::mutex> lock(m_mutex);
    std::uniform_int_distribution<int64_t> dist(min, max);
    return dist(m_gen);
}

double SecureRandom::GetDouble() {
    std::lock_guard<std::mutex> lock(m_mutex);
    std::uniform_real_distribution<double> dist(0.0, 1.0);
    return dist(m_gen);
}

//=============================================================================
// HASHING (Windows BCrypt)
//=============================================================================

#ifdef _WIN32

static std::string BytesToHex(const unsigned char* data, size_t len) {
    std::ostringstream oss;
    oss << std::hex << std::setfill('0');
    for (size_t i = 0; i < len; i++) {
        oss << std::setw(2) << static_cast<int>(data[i]);
    }
    return oss.str();
}

static LPCWSTR GetBCryptAlgorithm(HashAlgorithm algo) {
    switch (algo) {
        case HashAlgorithm::SHA256: return BCRYPT_SHA256_ALGORITHM;
        case HashAlgorithm::SHA512: return BCRYPT_SHA512_ALGORITHM;
        case HashAlgorithm::SHA1:   return BCRYPT_SHA1_ALGORITHM;
        case HashAlgorithm::MD5:    return BCRYPT_MD5_ALGORITHM;
        default:                    return BCRYPT_SHA256_ALGORITHM;
    }
}

std::string ComputeHash(HashAlgorithm algo, const std::string& data) {
    BCRYPT_ALG_HANDLE hAlg = nullptr;
    BCRYPT_HASH_HANDLE hHash = nullptr;
    NTSTATUS status;
    std::string result;

    status = BCryptOpenAlgorithmProvider(&hAlg, GetBCryptAlgorithm(algo), nullptr, 0);
    if (!BCRYPT_SUCCESS(status)) {
        return "";
    }

    DWORD hashLength = 0;
    DWORD resultLength = 0;
    status = BCryptGetProperty(hAlg, BCRYPT_HASH_LENGTH, (PUCHAR)&hashLength,
                                sizeof(hashLength), &resultLength, 0);
    if (!BCRYPT_SUCCESS(status)) {
        BCryptCloseAlgorithmProvider(hAlg, 0);
        return "";
    }

    std::vector<unsigned char> hashBuffer(hashLength);

    status = BCryptCreateHash(hAlg, &hHash, nullptr, 0, nullptr, 0, 0);
    if (!BCRYPT_SUCCESS(status)) {
        BCryptCloseAlgorithmProvider(hAlg, 0);
        return "";
    }

    status = BCryptHashData(hHash, (PUCHAR)data.data(), static_cast<ULONG>(data.size()), 0);
    if (!BCRYPT_SUCCESS(status)) {
        BCryptDestroyHash(hHash);
        BCryptCloseAlgorithmProvider(hAlg, 0);
        return "";
    }

    status = BCryptFinishHash(hHash, hashBuffer.data(), hashLength, 0);
    if (BCRYPT_SUCCESS(status)) {
        result = BytesToHex(hashBuffer.data(), hashLength);
    }

    BCryptDestroyHash(hHash);
    BCryptCloseAlgorithmProvider(hAlg, 0);

    return result;
}

std::string ComputeHMAC(HashAlgorithm algo, const std::string& key, const std::string& data) {
    BCRYPT_ALG_HANDLE hAlg = nullptr;
    BCRYPT_HASH_HANDLE hHash = nullptr;
    NTSTATUS status;
    std::string result;

    status = BCryptOpenAlgorithmProvider(&hAlg, GetBCryptAlgorithm(algo), nullptr,
                                          BCRYPT_ALG_HANDLE_HMAC_FLAG);
    if (!BCRYPT_SUCCESS(status)) {
        return "";
    }

    DWORD hashLength = 0;
    DWORD resultLength = 0;
    status = BCryptGetProperty(hAlg, BCRYPT_HASH_LENGTH, (PUCHAR)&hashLength,
                                sizeof(hashLength), &resultLength, 0);
    if (!BCRYPT_SUCCESS(status)) {
        BCryptCloseAlgorithmProvider(hAlg, 0);
        return "";
    }

    std::vector<unsigned char> hashBuffer(hashLength);

    status = BCryptCreateHash(hAlg, &hHash, nullptr, 0,
                               (PUCHAR)key.data(), static_cast<ULONG>(key.size()), 0);
    if (!BCRYPT_SUCCESS(status)) {
        BCryptCloseAlgorithmProvider(hAlg, 0);
        return "";
    }

    status = BCryptHashData(hHash, (PUCHAR)data.data(), static_cast<ULONG>(data.size()), 0);
    if (!BCRYPT_SUCCESS(status)) {
        BCryptDestroyHash(hHash);
        BCryptCloseAlgorithmProvider(hAlg, 0);
        return "";
    }

    status = BCryptFinishHash(hHash, hashBuffer.data(), hashLength, 0);
    if (BCRYPT_SUCCESS(status)) {
        result = BytesToHex(hashBuffer.data(), hashLength);
    }

    BCryptDestroyHash(hHash);
    BCryptCloseAlgorithmProvider(hAlg, 0);

    return result;
}

#elif defined(MOSIS_HAS_OPENSSL)

//=============================================================================
// HASHING (OpenSSL)
//=============================================================================

static std::string BytesToHex(const unsigned char* data, size_t len) {
    std::ostringstream oss;
    oss << std::hex << std::setfill('0');
    for (size_t i = 0; i < len; i++) {
        oss << std::setw(2) << static_cast<int>(data[i]);
    }
    return oss.str();
}

static const EVP_MD* GetOpenSSLAlgorithm(HashAlgorithm algo) {
    switch (algo) {
        case HashAlgorithm::SHA256: return EVP_sha256();
        case HashAlgorithm::SHA512: return EVP_sha512();
        case HashAlgorithm::SHA1:   return EVP_sha1();
        case HashAlgorithm::MD5:    return EVP_md5();
        default:                    return EVP_sha256();
    }
}

std::string ComputeHash(HashAlgorithm algo, const std::string& data) {
    const EVP_MD* md = GetOpenSSLAlgorithm(algo);
    unsigned char hash[EVP_MAX_MD_SIZE];
    unsigned int hash_len = 0;

    EVP_MD_CTX* ctx = EVP_MD_CTX_new();
    if (!ctx) return "";

    if (EVP_DigestInit_ex(ctx, md, nullptr) != 1 ||
        EVP_DigestUpdate(ctx, data.data(), data.size()) != 1 ||
        EVP_DigestFinal_ex(ctx, hash, &hash_len) != 1) {
        EVP_MD_CTX_free(ctx);
        return "";
    }

    EVP_MD_CTX_free(ctx);
    return BytesToHex(hash, hash_len);
}

std::string ComputeHMAC(HashAlgorithm algo, const std::string& key, const std::string& data) {
    const EVP_MD* md = GetOpenSSLAlgorithm(algo);
    unsigned char hmac_result[EVP_MAX_MD_SIZE];
    unsigned int hmac_len = 0;

    unsigned char* result = HMAC(md,
                                  key.data(), static_cast<int>(key.size()),
                                  reinterpret_cast<const unsigned char*>(data.data()),
                                  data.size(),
                                  hmac_result, &hmac_len);

    if (!result) return "";
    return BytesToHex(hmac_result, hmac_len);
}

#else

// Stub implementations when no crypto library is available
std::string ComputeHash(HashAlgorithm algo, const std::string& data) {
    (void)algo;
    (void)data;
    return "";
}

std::string ComputeHMAC(HashAlgorithm algo, const std::string& key, const std::string& data) {
    (void)algo;
    (void)key;
    (void)data;
    return "";
}

#endif

//=============================================================================
// LUA CRYPTO API
//=============================================================================

static const char* CRYPTO_RNG_KEY = "__mosis_crypto_rng";

static SecureRandom* GetRng(lua_State* L) {
    lua_getfield(L, LUA_REGISTRYINDEX, CRYPTO_RNG_KEY);
    if (lua_islightuserdata(L, -1)) {
        SecureRandom* rng = static_cast<SecureRandom*>(lua_touserdata(L, -1));
        lua_pop(L, 1);
        return rng;
    }
    lua_pop(L, 1);

    // Create a default RNG if none registered
    static SecureRandom default_rng;
    return &default_rng;
}

// crypto.randomBytes(n) -> string
static int lua_crypto_randomBytes(lua_State* L) {
    lua_Integer n = luaL_checkinteger(L, 1);

    if (n < 0) {
        return luaL_error(L, "crypto.randomBytes: count must be non-negative");
    }
    if (n > 1024) {
        return luaL_error(L, "crypto.randomBytes: count must not exceed 1024");
    }

    SecureRandom* rng = GetRng(L);
    std::string bytes = rng->GetBytes(static_cast<size_t>(n));

    lua_pushlstring(L, bytes.data(), bytes.size());
    return 1;
}

static HashAlgorithm ParseAlgorithm(const char* name) {
    if (strcmp(name, "sha256") == 0) return HashAlgorithm::SHA256;
    if (strcmp(name, "sha512") == 0) return HashAlgorithm::SHA512;
    if (strcmp(name, "sha1") == 0)   return HashAlgorithm::SHA1;
    if (strcmp(name, "md5") == 0)    return HashAlgorithm::MD5;
    return HashAlgorithm::SHA256;  // Default
}

// crypto.hash(algorithm, data) -> string
static int lua_crypto_hash(lua_State* L) {
    const char* algo_name = luaL_checkstring(L, 1);
    size_t data_len;
    const char* data = luaL_checklstring(L, 2, &data_len);

    // Limit input size
    if (data_len > 10 * 1024 * 1024) {
        return luaL_error(L, "crypto.hash: input too large (max 10MB)");
    }

    HashAlgorithm algo = ParseAlgorithm(algo_name);
    std::string result = ComputeHash(algo, std::string(data, data_len));

    if (result.empty()) {
        return luaL_error(L, "crypto.hash: failed to compute hash");
    }

    lua_pushstring(L, result.c_str());
    return 1;
}

// crypto.hmac(algorithm, key, data) -> string
static int lua_crypto_hmac(lua_State* L) {
    const char* algo_name = luaL_checkstring(L, 1);
    size_t key_len;
    const char* key = luaL_checklstring(L, 2, &key_len);
    size_t data_len;
    const char* data = luaL_checklstring(L, 3, &data_len);

    // Limit input sizes
    if (key_len > 1024) {
        return luaL_error(L, "crypto.hmac: key too large (max 1KB)");
    }
    if (data_len > 10 * 1024 * 1024) {
        return luaL_error(L, "crypto.hmac: data too large (max 10MB)");
    }

    HashAlgorithm algo = ParseAlgorithm(algo_name);
    std::string result = ComputeHMAC(algo, std::string(key, key_len),
                                      std::string(data, data_len));

    if (result.empty()) {
        return luaL_error(L, "crypto.hmac: failed to compute HMAC");
    }

    lua_pushstring(L, result.c_str());
    return 1;
}

// Helper to set a global in the real _G (bypassing any proxy)
static void SetGlobalInRealG(lua_State* L, const char* name) {
    // Stack: value to set as global

    // Get _G (might be a proxy)
    lua_rawgeti(L, LUA_REGISTRYINDEX, LUA_RIDX_GLOBALS);

    // Check if it has a metatable with __index (proxy pattern)
    if (lua_getmetatable(L, -1)) {
        lua_getfield(L, -1, "__index");
        if (lua_istable(L, -1)) {
            // Found real _G through proxy's __index
            // Stack: value, proxy, mt, real_G
            lua_pushvalue(L, -4);  // Copy value
            lua_setfield(L, -2, name);  // real_G[name] = value
            lua_pop(L, 4);  // pop real_G, mt, proxy, original value
            return;
        }
        lua_pop(L, 2);  // pop __index, metatable
    }

    // No proxy, set directly in _G
    // Stack: value, _G
    lua_pushvalue(L, -2);  // Copy value
    lua_setfield(L, -2, name);  // _G[name] = value
    lua_pop(L, 2);  // pop _G, original value
}

void RegisterCryptoAPI(lua_State* L) {
    // Create crypto table
    lua_newtable(L);

    lua_pushcfunction(L, lua_crypto_randomBytes);
    lua_setfield(L, -2, "randomBytes");

    lua_pushcfunction(L, lua_crypto_hash);
    lua_setfield(L, -2, "hash");

    lua_pushcfunction(L, lua_crypto_hmac);
    lua_setfield(L, -2, "hmac");

    // Set as global (bypassing proxy)
    SetGlobalInRealG(L, "crypto");
}

//=============================================================================
// SECURE MATH.RANDOM
//=============================================================================

static const char* MATH_RNG_KEY = "__mosis_math_rng";

// math.random([m [, n]]) - secure version
static int lua_secure_random(lua_State* L) {
    lua_getfield(L, LUA_REGISTRYINDEX, MATH_RNG_KEY);
    if (!lua_islightuserdata(L, -1)) {
        lua_pop(L, 1);
        return luaL_error(L, "math.random: RNG not initialized");
    }
    SecureRandom* rng = static_cast<SecureRandom*>(lua_touserdata(L, -1));
    lua_pop(L, 1);

    int nargs = lua_gettop(L);

    if (nargs == 0) {
        // Return double in [0.0, 1.0)
        lua_pushnumber(L, rng->GetDouble());
        return 1;
    } else if (nargs == 1) {
        // Return integer in [1, n]
        lua_Integer n = luaL_checkinteger(L, 1);
        if (n < 1) {
            return luaL_error(L, "math.random: interval is empty");
        }
        lua_pushinteger(L, rng->GetInt(1, n));
        return 1;
    } else {
        // Return integer in [m, n]
        lua_Integer m = luaL_checkinteger(L, 1);
        lua_Integer n = luaL_checkinteger(L, 2);
        if (m > n) {
            return luaL_error(L, "math.random: interval is empty");
        }
        lua_pushinteger(L, rng->GetInt(m, n));
        return 1;
    }
}

void RegisterSecureMathRandom(lua_State* L, SecureRandom* rng) {
    // Store RNG in registry
    lua_pushlightuserdata(L, rng);
    lua_setfield(L, LUA_REGISTRYINDEX, MATH_RNG_KEY);

    // Also store for crypto API
    lua_pushlightuserdata(L, rng);
    lua_setfield(L, LUA_REGISTRYINDEX, CRYPTO_RNG_KEY);

    // Get the math table
    lua_getglobal(L, "math");
    if (!lua_istable(L, -1)) {
        lua_pop(L, 1);
        return;
    }

    // Replace math.random with secure version
    lua_pushcfunction(L, lua_secure_random);
    lua_setfield(L, -2, "random");

    // Remove math.randomseed
    lua_pushnil(L);
    lua_setfield(L, -2, "randomseed");

    lua_pop(L, 1);  // Pop math table
}

} // namespace mosis