#include "blind.h"

#include "hash.h"
#include "primitives/transaction.h"
#include "random.h"
#include "util.h"

#include <secp256k1.h>
#include <secp256k1_rangeproof.h>
#include <secp256k1_surjectionproof.h>

static secp256k1_context* secp256k1_blind_context = NULL;

class Blind_ECC_Init {
public:
    Blind_ECC_Init() {
        assert(secp256k1_blind_context == NULL);

        secp256k1_context *ctx = secp256k1_context_create(SECP256K1_CONTEXT_SIGN | SECP256K1_CONTEXT_VERIFY);
        assert(ctx != NULL);

        secp256k1_blind_context = ctx;
    }

    ~Blind_ECC_Init() {
        secp256k1_context *ctx = secp256k1_blind_context;
        secp256k1_blind_context = NULL;

        if (ctx) {
            secp256k1_context_destroy(ctx);
        }
    }
};

static Blind_ECC_Init ecc_init_on_load;

bool UnblindOutput(const CKey &key, const CTxOut& txout, const CTxOutWitness& txoutwit, CAmount& amount_out, uint256& blinding_factor_out, CAsset& asset_out, uint256& asset_blinding_factor_out)
{
    if (!key.IsValid()) {
        return false;
    }
    CPubKey ephemeral_key(txout.nNonce.vchCommitment);
    if (!ephemeral_key.IsValid()) {
        return false;
    }
    uint256 nonce = key.ECDH(ephemeral_key);
    CSHA256().Write(nonce.begin(), 32).Finalize(nonce.begin());
    unsigned char msg[4096];
    size_t msg_size = 64;
    uint64_t min_value, max_value, amount;
    secp256k1_pedersen_commitment commit;
    if (!txout.nAsset.IsCommitment() || txout.nValue.IsExplicit())
        return false;

    secp256k1_generator gen;
    if (secp256k1_generator_parse(secp256k1_blind_context, &gen, &txout.nAsset.vchCommitment[0]) != 1)
        return false;
    if (secp256k1_pedersen_commitment_parse(secp256k1_blind_context, &commit, &txout.nValue.vchCommitment[0]) != 1)
        return false;
    int res = secp256k1_rangeproof_rewind(secp256k1_blind_context, blinding_factor_out.begin(), &amount, msg, &msg_size, nonce.begin(), &min_value, &max_value, &commit, &txoutwit.vchRangeproof[0], txoutwit.vchRangeproof.size(), txout.scriptPubKey.size()? &txout.scriptPubKey.front(): NULL, txout.scriptPubKey.size(), &gen);
    secp256k1_generator recoveredGen;

    if (!res || amount > (uint64_t)MAX_MONEY || !MoneyRange((CAmount)amount) || msg_size != 64 || secp256k1_generator_generate_blinded(secp256k1_blind_context, &recoveredGen, msg+32, msg+64) != 1 || !memcmp(&gen, &recoveredGen, 33)) {
        amount_out = 0;
        blinding_factor_out = uint256();
        asset_out.SetNull();
        asset_blinding_factor_out = uint256();
        return false;
    } else {
        amount_out = (CAmount)amount;
        asset_out = CAsset(std::vector<unsigned char>(msg, msg+32));
        asset_blinding_factor_out = uint256(std::vector<unsigned char>(msg+32, msg+64));
        return true;
    }
}

int BlindOutputs(std::vector<uint256 >& input_blinding_factors, const std::vector<uint256 >& input_asset_blinding_factors, const std::vector<CAsset >& input_assets, const std::vector<CAmount >& input_amounts, std::vector<uint256 >& output_blinding_factors, std::vector<uint256 >& output_asset_blinding_factors, const std::vector<CPubKey>& output_pubkeys, CMutableTransaction& tx, std::vector<std::vector<unsigned char> >* auxiliary_generators)
{
    // Sanity check input data and output_pubkey size, clear other output data
    assert(tx.vout.size() >= output_pubkeys.size());
    output_blinding_factors.clear();
    output_blinding_factors.resize(tx.vout.size());
    output_asset_blinding_factors.clear();
    output_asset_blinding_factors.resize(tx.vout.size());
    assert(tx.vin.size() == input_blinding_factors.size());
    assert(tx.vin.size() == input_asset_blinding_factors.size());
    assert(tx.vin.size() == input_assets.size());
    assert(tx.vin.size() == input_amounts.size());

    std::vector<unsigned char*> blindptrs;
    std::vector<const unsigned char*> assetblindptrs;
    std::vector<uint64_t> blindedAmounts;
    blindptrs.reserve(tx.vout.size() + tx.vin.size());
    assetblindptrs.reserve(tx.vout.size() + tx.vin.size());

    int ret;
    int nBlindAttempts = 0, nSuccessfullyBlinded = 0;

    //Surjection proof prep

    // Needed to surj init, only matches to output asset matters, rest can be garbage
    std::vector<secp256k1_fixed_asset_tag> inputAssets;

    // Needed to construct the proof itself. Generators must match final transaction to be valid
    std::vector<secp256k1_generator> inputAssetGenerators;
    inputAssets.resize(tx.vin.size());
    inputAssetGenerators.resize(tx.vin.size());
    for (size_t i = 0; i < tx.vin.size(); i++) {
        // If non-empty generator exists, parse
        if (auxiliary_generators && auxiliary_generators->size() > i && auxiliary_generators[i].size() == 33) {
            // Parse generator here
            ret = secp256k1_generator_parse(secp256k1_blind_context, &inputAssetGenerators[i], &(*auxiliary_generators)[i][0]);
            if (ret != 1) {
                return -1;
            }
        } else {
            // Needs to be non-null
            if (input_assets[i].IsNull()) {
                return -1;
            }
            ret = secp256k1_generator_generate_blinded(secp256k1_blind_context, &inputAssetGenerators[i], input_assets[i].begin(), input_asset_blinding_factors[i].begin());
            assert(ret == 1);
        }
        memcpy(&inputAssets[i], input_assets[i].begin(), 32);
    }

    //Total blinded inputs that you own (that you are balancing against)
    int nBlindsIn = 0;
    for (size_t nIn = 0; nIn < tx.vin.size(); nIn++) {
        if (!input_blinding_factors[nIn].IsNull() || !input_asset_blinding_factors[nIn].IsNull()) {
            if (input_amounts[nIn] < 0) {
                return -1;
            }
            blindptrs.push_back(input_blinding_factors[nIn].begin());
            assetblindptrs.push_back(input_asset_blinding_factors[nIn].begin());
            blindedAmounts.push_back(input_amounts[nIn]);
            nBlindsIn++;
        }
    }

    //Number of outputs to blind
    int nToBlind = 0;
    for (size_t nOut = 0; nOut < output_pubkeys.size(); nOut++) {
        if (output_pubkeys[nOut].IsValid()) {
            // Keys must be valid and outputs completely unblinded or else call fails
            if (!output_pubkeys[nOut].IsFullyValid() ||
                (!tx.vout[nOut].nValue.IsExplicit() || !tx.vout[nOut].nAsset.IsExplicit()) ||
                    tx.vout[nOut].IsFee()) {
                return -1;
            }
            nToBlind++;
         }
    }


    //Running total of newly blinded outputs
    static const unsigned char diff_zero[32] = {0};
    unsigned char blind[tx.vout.size()][32];
    unsigned char asset_blind[tx.vout.size()][32];
    secp256k1_pedersen_commitment commit;
    secp256k1_generator gen;
    CAsset assetID;

    // This section of code *only* deals with unblinded outputs
    // that we want to blind
    for (size_t nOut = 0; nOut < output_pubkeys.size(); nOut++) {
        CTxOut& out = tx.vout[nOut];
        if (out.nValue.IsExplicit() && output_pubkeys[nOut].IsFullyValid()) {
            nBlindAttempts++;
            CConfidentialAsset& asset = out.nAsset;
            CConfidentialValue& value = out.nValue;
            CAmount amount = value.GetAmount();
            assetID = out.nAsset.GetAsset();
            blindedAmounts.push_back(value.GetAmount());

            GetRandBytes(&blind[nBlindAttempts-1][0], 32);
            GetRandBytes(&asset_blind[nBlindAttempts-1][0], 32);
            blindptrs.push_back(&blind[nBlindAttempts-1][0]);
            assetblindptrs.push_back(&asset_blind[nBlindAttempts-1][0]);

            // Last blinding factor r' is set as -(output's (vr + r') - input's (vr + r')).
            // Before modifying the transaction or return arguments we must
            // ensure the final blinding factor to not be its corresponding -vr (aka unblinded),
            // or 0, in the case of 0-value output, insisting on additional output to blind.
            if (nBlindAttempts == nToBlind) {

                // Can't successfully blind in this case, since -vr = r
                // This check is assuming blinds are generated randomly
                // Adversary would need to create all input blinds
                // therefore would already know all your summed output amount anyways.
                if (nBlindAttempts == 1 && nBlindsIn == 0) {
                    return nSuccessfullyBlinded;
                }

                // Generate value we intend to insert
                ret = secp256k1_pedersen_blind_generator_blind_sum(secp256k1_blind_context, &blindedAmounts[0], &assetblindptrs[0], &blindptrs[0], nBlindAttempts + nBlindsIn, nBlindsIn);
                assert(ret);

                // Resulting blinding factor shouldn't be 0
                if (memcmp(diff_zero, &blind[nBlindAttempts-1][0], 32) == 0) {
                   return nSuccessfullyBlinded;
                }
            }

            if (tx.wit.vtxoutwit.size() <= nOut) {
                tx.wit.vtxoutwit.resize(tx.vout.size());
            }
            CTxOutWitness& txoutwit = tx.wit.vtxoutwit[nOut];

            output_blinding_factors[nOut] = uint256(std::vector<unsigned char>(blindptrs[blindptrs.size()-1], blindptrs[blindptrs.size()-1]+32));
            output_asset_blinding_factors[nOut] = uint256(std::vector<unsigned char>(assetblindptrs[assetblindptrs.size()-1], assetblindptrs[assetblindptrs.size()-1]+32));

            //Blind the asset ID
            ret = secp256k1_generator_generate_blinded(secp256k1_blind_context, &gen, assetID.begin(), assetblindptrs[assetblindptrs.size()-1]);
            assert(ret == 1);
            ret = secp256k1_generator_serialize(secp256k1_blind_context, &asset.vchCommitment[0], &gen);
            assert(ret != 0);

            // Create value commitment
            value.vchCommitment.resize(CConfidentialValue::nCommittedSize);
            ret = secp256k1_pedersen_commit(secp256k1_blind_context, &commit, (unsigned char*)blindptrs.back(), amount, &gen);
            assert(ret != 0);
            secp256k1_pedersen_commitment_serialize(secp256k1_blind_context, &value.vchCommitment[0], &commit);
            assert(value.IsValid());

            // Generate ephemeral key for ECDH nonce generation
            CKey ephemeral_key;
            ephemeral_key.MakeNewKey(true);
            CPubKey ephemeral_pubkey = ephemeral_key.GetPubKey();
            assert(ephemeral_pubkey.size() == CConfidentialNonce::nCommittedSize);
            out.nNonce.vchCommitment.resize(ephemeral_pubkey.size());
            memcpy(&out.nNonce.vchCommitment[0], &ephemeral_pubkey[0], ephemeral_pubkey.size());
            // Generate nonce
            uint256 nonce = ephemeral_key.ECDH(output_pubkeys[nOut]);
            CSHA256().Write(nonce.begin(), 32).Finalize(nonce.begin());

            // Prep range proof
            size_t nRangeProofLen = 5134;
            // TODO: smarter min_value selection
            txoutwit.vchRangeproof.resize(nRangeProofLen);

            // Compose sidechannel message to convey asset info (ID and asset blinds)
            unsigned char assetsMessage[64];
            memcpy(assetsMessage, assetID.begin(), 32);
            memcpy(assetsMessage+32,  assetblindptrs[assetblindptrs.size()-1], 32);

            // Sign rangeproof
            int res = secp256k1_rangeproof_sign(secp256k1_blind_context, &txoutwit.vchRangeproof[0], &nRangeProofLen, 0, &commit, blindptrs.back(), nonce.begin(), std::min(std::max((int)GetArg("-ct_exponent", 0), -1),18), std::min(std::max((int)GetArg("-ct_bits", 32), 1), 51), amount, assetsMessage, sizeof(assetsMessage), out.scriptPubKey.size() ? &out.scriptPubKey.front() : NULL, out.scriptPubKey.size(), &gen);
            txoutwit.vchRangeproof.resize(nRangeProofLen);
            // TODO: do something smarter here
            assert(res);

            // Create surjection proof
            size_t nInputsToSelect = std::min((size_t)3, input_assets.size());
            unsigned char randseed[32];
            GetRandBytes(randseed, 32);
            size_t input_index;
            secp256k1_surjectionproof proof;
            secp256k1_fixed_asset_tag tag;
            memcpy(&tag, assetID.begin(), 32);
            if (secp256k1_surjectionproof_initialize(secp256k1_blind_context, &proof, &input_index, &inputAssets[0], input_assets.size(), nInputsToSelect, &tag, 100, randseed) == 0) {
                continue;
            }
            ret = secp256k1_surjectionproof_generate(secp256k1_blind_context, &proof, &inputAssetGenerators[0], inputAssetGenerators.size(), &gen, input_index, input_asset_blinding_factors[input_index].begin(), assetblindptrs[assetblindptrs.size()-1]);
            assert(ret == 1);
            ret = secp256k1_surjectionproof_verify(secp256k1_blind_context, &proof, &inputAssetGenerators[0], inputAssetGenerators.size(), &gen);
            assert(ret != 0);

            size_t output_len = secp256k1_surjectionproof_serialized_size(secp256k1_blind_context, &proof);
            txoutwit.vchSurjectionproof.resize(output_len);
            secp256k1_surjectionproof_serialize(secp256k1_blind_context, &txoutwit.vchSurjectionproof[0], &output_len, &proof);

            // Successfully blinded this output
            nSuccessfullyBlinded++;
        }
    }

    return nSuccessfullyBlinded;
}