#include <pybind11/pybind11.h>
#include <pybind11/stl.h>  // 支持 std::vector, std::map 等
#include <nlohmann/json.hpp>
#include <cstring>
#include <cstdint>
#include <vector>
#include <string>
#include <fstream>
#include <array>
#include <openssl/evp.h>
#include <algorithm>
#include "native_logger.hpp"

namespace netcore{

    // OTA AEAD format: MAGIC(7) | nonce(12) | ciphertext(N) | tag(16)
    constexpr const char* kOtaMagic = "AROTAE1";
    constexpr size_t kOtaMagicLen = 7;
    constexpr size_t kGcmNonceLen = 12;
    constexpr size_t kGcmTagLen = 16;

    // 固定 32-byte AES-256-GCM key（提高被直接查看的成本；不是绝对安全）
    // 注意：需要与打包端传入的 --aead-key-hex 保持一致。
    static std::array<uint8_t, 32> ota_key_bytes() {
        // 简单拆分混淆：key = a XOR b
        static const std::array<uint8_t, 32> a = {
            0x92,0x99,0x4d,0x06,0x6f,0xb6,0xa6,0x3d,0x85,0x08,0xbe,0x73,0x5e,0x73,0x4d,0x8a,
            0x53,0x88,0xe6,0x99,0xfc,0x10,0x29,0xb9,0x16,0x9b,0xe7,0x0c,0x65,0x21,0x1c,0xce
        };
        static const std::array<uint8_t, 32> b = {
            0xcf,0x60,0xa2,0xc2,0x32,0x7a,0x61,0xb0,0x4c,0x8e,0x8a,0x62,0x31,0xc7,0x82,0xff,
            0xec,0xac,0xa1,0x04,0x2a,0x4d,0xaa,0xf2,0xb0,0x5b,0x39,0x2b,0xf4,0xb3,0xad,0xad
        };
        std::array<uint8_t, 32> k{};
        for (size_t i = 0; i < k.size(); i++) k[i] = static_cast<uint8_t>(a[i] ^ b[i]);
        return k;
    }

    static bool read_file_all(const std::string& path, std::vector<uint8_t>& out) {
        std::ifstream ifs(path, std::ios::binary);
        if (!ifs) return false;
        ifs.seekg(0, std::ios::end);
        std::streampos size = ifs.tellg();
        if (size <= 0) return false;
        ifs.seekg(0, std::ios::beg);
        out.resize(static_cast<size_t>(size));
        if (!ifs.read(reinterpret_cast<char*>(out.data()), size)) return false;
        return true;
    }

    static bool write_file_all(const std::string& path, const uint8_t* data, size_t len) {
        std::ofstream ofs(path, std::ios::binary | std::ios::trunc);
        if (!ofs) return false;
        ofs.write(reinterpret_cast<const char*>(data), static_cast<std::streamsize>(len));
        return static_cast<bool>(ofs);
    }

    bool decrypt_ota_file_impl(const std::string& input_path, const std::string& output_zip_path) {
        std::vector<uint8_t> in;
        if (!netcore::read_file_all(input_path, in)) {
            netcore::log_error(std::string("decrypt_ota_file: read failed: ") + input_path);
            return false;
        }

        const size_t min_len = kOtaMagicLen + kGcmNonceLen + kGcmTagLen + 1;
        if (in.size() < min_len) {
            netcore::log_error("decrypt_ota_file: too short");
            return false;
        }
        if (!std::equal(in.begin(), in.begin() + kOtaMagicLen, reinterpret_cast<const uint8_t*>(kOtaMagic))) {
            netcore::log_error("decrypt_ota_file: bad magic");
            return false;
        }

        const uint8_t* nonce = in.data() + kOtaMagicLen;
        const uint8_t* ct_and_tag = in.data() + kOtaMagicLen + kGcmNonceLen;
        const size_t ct_and_tag_len = in.size() - (kOtaMagicLen + kGcmNonceLen);
        if (ct_and_tag_len <= kGcmTagLen) {
            netcore::log_error("decrypt_ota_file: no ciphertext");
            return false;
        }
        const size_t ciphertext_len = ct_and_tag_len - kGcmTagLen;
        const uint8_t* ciphertext = ct_and_tag;
        const uint8_t* tag = ct_and_tag + ciphertext_len;

        std::vector<uint8_t> plain(ciphertext_len);
        int out_len1 = 0;
        int out_len2 = 0;

        EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
        if (!ctx) {
            netcore::log_error("decrypt_ota_file: EVP_CIPHER_CTX_new failed");
            return false;
        }

        bool ok = false;
        auto key = ota_key_bytes();

        do {
            if (1 != EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr)) {
                netcore::log_error("decrypt_ota_file: DecryptInit failed");
                break;
            }
            if (1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, static_cast<int>(kGcmNonceLen), nullptr)) {
                netcore::log_error("decrypt_ota_file: set ivlen failed");
                break;
            }
            if (1 != EVP_DecryptInit_ex(ctx, nullptr, nullptr, key.data(), nonce)) {
                netcore::log_error("decrypt_ota_file: set key/iv failed");
                break;
            }
            if (1 != EVP_DecryptUpdate(ctx, plain.data(), &out_len1, ciphertext, static_cast<int>(ciphertext_len))) {
                netcore::log_error("decrypt_ota_file: update failed");
                break;
            }
            if (1 != EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, static_cast<int>(kGcmTagLen), const_cast<uint8_t*>(tag))) {
                netcore::log_error("decrypt_ota_file: set tag failed");
                break;
            }
            if (1 != EVP_DecryptFinal_ex(ctx, plain.data() + out_len1, &out_len2)) {
                netcore::log_error("decrypt_ota_file: final failed (auth tag mismatch?)");
                break;
            }
            const size_t plain_len = static_cast<size_t>(out_len1 + out_len2);
            if (!netcore::write_file_all(output_zip_path, plain.data(), plain_len)) {
                netcore::log_error(std::string("decrypt_ota_file: write failed: ") + output_zip_path);
                break;
            }
            ok = true;
        } while (false);

        EVP_CIPHER_CTX_free(ctx);
        return ok;
    }
}