#include "StateIntegrity.hpp"
#include "Base64.h"

extern "C" {
#include "ed25519.h"
}

using namespace juce;

std::string StateIntegrity::generateSeed() {
    std::string seed(SEED_SIZE, 0);
    ed25519_create_seed((unsigned char*)seed.data());
    return seed;
}

StateIntegrity::StateIntegrity(const std::string_view & seed) {
    if(seed.size() != SEED_SIZE) {
        throw std::runtime_error("Wrong seed size");
    }
    
    ed25519_create_keypair(publicKey.data(), privateKey.data(), (unsigned char*)seed.data());
}

StateIntegrity::~StateIntegrity() {
    privateKey.fill(0);
}

std::string StateIntegrity::sign(const std::string_view & data) {
    std::array<unsigned char, SIGNATURE_SIZE> signature;
    ed25519_sign(signature.data(), (unsigned char*)data.data(), data.size(), publicKey.data(), privateKey.data());
    // output = publicKey.signature.data
    std::string output(data.size() + PUBLIC_KEY_SIZE + SIGNATURE_SIZE, 0);
    std::copy(publicKey.begin(), publicKey.end(), output.begin());
    std::copy(signature.begin(), signature.end(), output.begin() + PUBLIC_KEY_SIZE);
    std::copy(data.begin(), data.end(), output.begin() + PUBLIC_KEY_SIZE + SIGNATURE_SIZE);
    return output;
}

std::string StateIntegrity::verify(const std::string_view & data) {
    // data = publicKey.signature.data
    if(data.size() < PUBLIC_KEY_SIZE + SIGNATURE_SIZE) {
        throw SignatureVerificationFailed(Private());
    }
    public_key_t messagePublicKey;
    std::copy(data.begin(), data.begin() + PUBLIC_KEY_SIZE, messagePublicKey.begin());
    std::array<unsigned char, SIGNATURE_SIZE> messageSignature;
    std::copy(data.begin() + PUBLIC_KEY_SIZE, data.begin() + PUBLIC_KEY_SIZE + SIGNATURE_SIZE, messageSignature.begin());
    std::string messageRest(data.begin() + PUBLIC_KEY_SIZE + SIGNATURE_SIZE, data.end());
    
    int verifyResult = ed25519_verify(messageSignature.data(),
                                      (unsigned char*)messageRest.data(),
                                      messageRest.size(),
                                      messagePublicKey.data());
    if(verifyResult == 0) {
        throw SignatureVerificationFailed(Private());
    }
    
    if (messagePublicKey == publicKey) {
        return messageRest;
    }
    
    auto trustRecord = trustStore.find(messagePublicKey);
    if(trustRecord == trustStore.end() || !trustRecord->second) {
        throw PublicKeyNotTrusted(Private(), std::string(messagePublicKey.begin(), messagePublicKey.end()), trustRecord != trustStore.end());
    }
    
    return messageRest;
}

void StateIntegrity::updateKeyTrust(const std::string_view & key, std::optional<bool> trust) {
    if(key.size() != PUBLIC_KEY_SIZE) {
        throw std::runtime_error("Key has wrong length");
    }
    public_key_t targetKey;
    std::copy(key.begin(), key.end(), targetKey.begin());
    
    if(trust.has_value()) {
        trustStore.insert(std::make_pair(targetKey, trust.value()));
    } else {
        auto it = trustStore.find(targetKey);
        if(it != trustStore.end()) {
            trustStore.erase(it);
        }
    }
}

std::optional<bool> StateIntegrity::getKeyTrust(const std::string_view & key) {
    if(key.size() != PUBLIC_KEY_SIZE) {
        throw std::runtime_error("Key has wrong length");
    }
    public_key_t targetKey;
    std::copy(key.begin(), key.end(), targetKey.begin());
    
    auto it = trustStore.find(targetKey);
    if(it == trustStore.end()) {
        return {};
    }
    
    return it->second;
}

void StateIntegrity::loadTrustStore(const juce::XmlElement & elem) {
    std::map<public_key_t, bool> newTrustStore;
    
    if(elem.getTagName() != "TrustStore") {
        throw std::runtime_error("Malformed trust store");
    }
    for(auto child : elem.getChildIterator()) {
        if(child->getTagName() != "PublicKey") {
            throw std::runtime_error("Malformed entry in trust store");
        }
        auto elemPublicKeyStr = convertFromBase64(child->getStringAttribute("value"));
        if(elemPublicKeyStr.size() != PUBLIC_KEY_SIZE) {
            throw std::runtime_error("Trust store public key has wrong length");
        }
        public_key_t elemPublicKey;
        std::copy(elemPublicKeyStr.begin(), elemPublicKeyStr.end(), elemPublicKey.begin());
        auto elemTrust = child->getBoolAttribute("trust");
        
        newTrustStore.insert(std::make_pair(elemPublicKey, elemTrust));
    }
    
    trustStore.clear();
    trustStore = std::move(newTrustStore);
}

std::unique_ptr<juce::XmlElement> StateIntegrity::saveTrustStore() {
    auto elem = std::make_unique<XmlElement>("TrustStore");
    for(auto & trustStoreEntry : trustStore) {
        auto subelem = std::make_unique<XmlElement>("PublicKey");
        auto & key = trustStoreEntry.first;
        subelem->setAttribute("value", convertToBase64(std::string_view((char*)key.data(), key.size())));
        subelem->setAttribute("trust", trustStoreEntry.second);
        elem->addChildElement(subelem.release());
    }
    return elem;
}

void StateIntegrity::clearTrustStore() {
    trustStore.clear();
}