#include <fstream>
#include <vector>
#include <iostream>
#include <cstring>
#include "keya.h"

using namespace std;

// Helper to pad data to 16 bytes
void pad_data(vector<unsigned char>& data)
{
    size_t pad_len = 16 - (data.size() % 16);
    if (pad_len < 16) {
        data.insert(data.end(), pad_len, 0); // simple zero padding
    }
}

// Read file into vector
vector<unsigned char> read_file(const string& filename)
{
    ifstream file(filename, ios::binary);
    return vector<unsigned char>((istreambuf_iterator<char>(file)), istreambuf_iterator<char>());
}

// Write vector into file
void write_file(const string& filename, const vector<unsigned char>& data)
{
    ofstream file(filename, ios::binary);
    file.write(reinterpret_cast<const char*>(data.data()), data.size());
}

// Encrypt buffer using KeyA
int encrypt_buffer(vector<unsigned char>& data, unsigned char* kp, unsigned char* kr, unsigned char* kw)
{
    KeyA keya;
    int result = keya.OpenDevice(0);
    if (result != 0) return result;

    result = keya.Init(kp, kr, kw, NULL);
    if (result != 0) {
        keya.CloseDevice();
        return result;
    }

    unsigned char temp[16];
    for (size_t i = 0; i < data.size(); i += 16) {
        memcpy(temp, &data[i], 16);
        result = keya.EncByKV(0, temp, temp);
        if (result != 0) {
            keya.CloseDevice();
            return result;
        }
        memcpy(&data[i], temp, 16);
    }

    keya.CloseDevice();
    return 0;
}

// Decrypt buffer using KeyA
int decrypt_buffer(vector<unsigned char>& data, unsigned char* kp, unsigned char* kr, unsigned char* kw)
{
    KeyA keya;
    int result = keya.OpenDevice(0);
    if (result != 0) return result;

    result = keya.Init(kp, kr, kw, NULL);
    if (result != 0) {
        keya.CloseDevice();
        return result;
    }

    unsigned char temp[16];
    for (size_t i = 0; i < data.size(); i += 16) {
        memcpy(temp, &data[i], 16);
        result = keya.DecByKV(0, temp, temp);
        if (result != 0) {
            keya.CloseDevice();
            return result;
        }
        memcpy(&data[i], temp, 16);
    }

    keya.CloseDevice();
    return 0;
}

// Save size as first 8 bytes (uint64_t) of encrypted file
void prepend_size(vector<unsigned char>& data, uint64_t size)
{
    vector<unsigned char> result(8);
    for (int i = 0; i < 8; ++i) {
        result[i] = (size >> (i * 8)) & 0xFF;
    }
    result.insert(result.end(), data.begin(), data.end());
    data = result;
}

// Extract original size from first 8 bytes
uint64_t extract_size(vector<unsigned char>& data)
{
    uint64_t size = 0;
    for (int i = 7; i >= 0; --i) {
        size = (size << 8) | data[i];
    }
    data.erase(data.begin(), data.begin() + 8);
    return size;
}

int main()
{
    // Your keys
    unsigned char g_Kp[16]  = {0x74,0xD8,0xCD,0x0F,0xC6,0x6A,0x09,0x69,0xB1,0x4A,0xC5,0xD1,0xE9,0x7A,0x9D,0x07};
    unsigned char g_Kr[16]  = {0xDA,0x58,0x86,0x3C,0xD6,0x65,0xDF,0xE2,0x84,0x04,0xD3,0x4E,0x74,0xEB,0xBD,0x16};
    unsigned char g_Kw[16]  = {0xC8,0xA8,0x16,0x3C,0x59,0xEF,0x9C,0xC4,0xC0,0xEC,0x48,0x41,0x4B,0xE5,0x11,0x37};

    unsigned char* kp = g_Kp;
    unsigned char* kr = g_Kr;
    unsigned char* kw = g_Kw;

    // Step 1: Encrypt JSON file
    vector<unsigned char> json_data = read_file("orig.txt");

    uint64_t original_size = json_data.size(); // Save original size
    pad_data(json_data); // pad to multiple of 16

    if (encrypt_buffer(json_data, kp, kr, kw) == 0) {
        prepend_size(json_data, original_size); // Put size in front
        write_file("data.enc", json_data);
        cout << "Encryption successful. Encrypted file saved as 'data.enc'." << endl;
    } else {
        cout << "Encryption failed." << endl;
    }

    // Step 2: Decrypt encrypted file
    vector<unsigned char> enc_data = read_file("hashes.enc");
    uint64_t saved_size = extract_size(enc_data); // remove first 8 bytes, get size

    if (decrypt_buffer(enc_data, kp, kr, kw) == 0) {
        enc_data.resize(saved_size); // trim padding
        write_file("hash.txt", enc_data);
        cout << "Decryption successful. Decrypted file saved as 'data_decrypted.txt'." << endl;
    } else {
        cout << "Decryption failed." << endl;
    }

    return 0;
}