#include "kms/kms_interface.hpp"
#include "kms/kms_core_local.hpp"
#include "en_config.h"
#include <iostream>
#include <string>
#include <fstream>
#include <chrono>
#include <vector>


using namespace std;

// 默认数据库名
const string DEFAULT_DBNAME = "default_db";

void print_usage() {
    cout << "Usage: KMS <command> <username> [args...] [--time]\n";
    cout << "Commands:\n";
    cout << "  hasCmk <username>\n";
    cout << "  createCmk <username> <rotate_period>\n";
    cout << "  deleteCmk <username>\n";
    cout << "  describeCmk <username> [decrypt]\n";
    cout << "  rotateCmkNow <username>\n";
    cout << "  CMK_auto_rotate <username> on/off/status\n";
    cout << "  encryptData <username> <data>\n";
    cout << "  decryptData <username> <data>\n";
    cout << "  createDek <username> <column_name>\n";
}

bool init_and_get_instance(const string& username) {
    LocalKmsCore::registerLocalKms();

    LocalKmsConfig config(
        CMK_FILE_PATH,
        CMK_KEY_PATH,
        CMK_ROTATE_STATUS_PATH,
        username.c_str(),
        DEFAULT_DBNAME.c_str()
    );

    if (!KMSInterface::initialize(config)) {
        cerr << "ERROR: KMS initialization failed\n";
        return false;
    }

    if (!KMSInterface::getInstance()) {
        cerr << "ERROR: KMS instance not available\n";
        return false;
    }

    return true;
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        print_usage();
        return 1;
    }

    // ======== 解析参数，处理 --time =========
    bool measure_time = false;
    vector<string> args;
    args.reserve(argc);

    for (int i = 1; i < argc; ++i) {
        if (string(argv[i]) == "--time") {
            measure_time = true;
        } else {
            args.push_back(argv[i]); // 保留非 --time 参数
        }
    }

    if (args.size() < 2) {
        print_usage();
        return 1;
    }

    string cmd = args[0];
    string username = args[1];

    if (cmd == "--help" || cmd == "-h") {
        print_usage();
        return 0;
    }

    // ======== 初始化 KMS =========
    if (!init_and_get_instance(username)) {
        return 1;
    }

    bool ok = false;
    string result;

    // ======== 计时开始 =========
    auto start = std::chrono::high_resolution_clock::now();

    // ========= KMS 命令处理 =========
    if (cmd == "hasCmk") {
        ok = KMSInterface::hasCmk(username);
        cout << (ok ? "true" : "false") << "\n";
    }
    else if (cmd == "createCmk" && args.size() == 3) {
        int rotate_period = stoi(args[2]);
        ok = KMSInterface::createCmk(username, rotate_period);
        cout << (ok ? "OK" : "ERROR") << "\n";
    }
    else if (cmd == "deleteCmk") {
        ok = KMSInterface::deleteCmk(username);
        cout << (ok ? "OK" : "ERROR") << "\n";
    }
    else if (cmd == "describeCmk") {
        bool decrypt = (args.size() == 3 && args[2] == "decrypt");
        ok = KMSInterface::describeCmk(username, result, decrypt);
        cout << (ok ? result : "ERROR") << "\n";
    }
    else if (cmd == "rotateCmkNow") {
        ok = KMSInterface::rotateCmkNow(username);
        cout << (ok ? "OK" : "ERROR") << "\n";
    }
    else if (cmd == "CMK_auto_rotate" && args.size() == 3) {
        string action = args[2];
        ok = KMSInterface::handleAutoRotateCmd(&action, username);
        cout << (ok ? action : "ERROR") << "\n";
    }
    else if (cmd == "encryptData" && args.size() == 3) {
        string input = args[2];
        ifstream infile(input);
        if (infile.good()) {
            string output_file = input + ".enc";
            ofstream outfile(output_file);
            string line;
            while (getline(infile, line)) {
                if (line.empty()) continue;
                string line_result = line;
                ok = KMSInterface::encryptData(line_result);
                outfile << (ok ? line_result : "ERROR") << "\n";
            }
            cout << "Encrypted result written to: " << output_file << "\n";
        } else {
            string str_result = input;
            ok = KMSInterface::encryptData(str_result);
            cout << (ok ? str_result : "ERROR") << "\n";
        }
    }
    else if (cmd == "decryptData" && args.size() == 3) {
        string input = args[2];
        ifstream infile(input);
        if (infile.good()) {
            string output_file = input + ".dec";
            ofstream outfile(output_file);
            string line;
            while (getline(infile, line)) {
                string line_result = line;
                ok = KMSInterface::decryptData(line_result);
                outfile << (ok ? line_result : "ERROR") << "\n";
            }
            cout << "Decryption results written to " << output_file << "\n";
        } else {
            string str_result = input;
            ok = KMSInterface::decryptData(str_result);
            cout << (ok ? str_result : "ERROR") << "\n";
        }
    }
    else if (cmd == "createDek" && args.size() == 3) {
        string input = args[2];
        ifstream infile(input);
        if (infile.good()) {
            string output_file = input + ".dek";
            ofstream outfile(output_file);
            string line;
            while (getline(infile, line)) {
                if (line.empty()) continue;
                string col_dek;
                ok = KMSInterface::createDek(col_dek, line);
                outfile << line << " : " << (ok ? col_dek : "ERROR") << "\n";
            }
            cout << "DEK creation results written to " << output_file << "\n";
        } else {
            string col_dek;
            ok = KMSInterface::createDek(col_dek, input);
            cout << (ok ? (input + " : " + col_dek) : "ERROR") << "\n";
        }
    }
    else {
        print_usage();
        return 1;
    }

    // ======== 计时结束 =========
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = chrono::duration_cast<chrono::microseconds>(end - start);

    if (measure_time) {
        cout << "[Time used]: " << duration.count() << " us\n";
    }

    return ok ? 0 : 1;
}
