/*
 * This file is part of the Flowee project
 * Copyright (C) 2019-2022 Tom Zander <tom@flowee.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "TransactionBuilder.h"

#include <hash.h>
#include <primitives/PublicKey.h>
#include <primitives/transaction.h>

class TransactionBuilderPrivate
{
public:
    void checkCurInput()
    {
        assert(0 <= curInput);
        assert(transaction.vin.size() > size_t(curInput));
        if (0 > curInput || transaction.vin.size() <= size_t(curInput))
            throw std::runtime_error("current input out of range");
    }
    void checkCurOutput()
    {
        assert(curOutput >= 0);
        assert(size_t(curOutput) < transaction.vout.size());
    }

    CMutableTransaction transaction;

    TransactionBuilder::LockingOptions defaultLocking = TransactionBuilder::NoLocking;
    int curInput = -1, curOutput = -1;

    struct SignInfo {
        uint8_t hashType = 0;
        int64_t amount = 0;
        PrivateKey privKey;
        CScript prevOutScript;
        TransactionBuilder::SignatureType signatureType = TransactionBuilder::Schnorr;
    };
    std::vector<SignInfo> signInfo;

    bool anonimize = false;
};

namespace {

// helper struct to sort index-of-inputs
struct InputSorter {
    const CMutableTransaction &tx;
    explicit InputSorter(const CMutableTransaction &transaction) : tx(transaction) {}
    bool operator()(int a, int b) const {
        assert(a >= 0);
        assert(b >= 0);
        assert((int)tx.vin.size() > a);
        assert((int)tx.vin.size() > b);
        const COutPoint &poa = tx.vin.at(a).prevout;
        const COutPoint &pob = tx.vin.at(b).prevout;

        auto rc = poa.hash.Compare(pob.hash);
        if (rc < 0)
            return true;
        if (rc > 0)
            return false;
        // equal, check output-index
        return poa.n < pob.n;
    }
};

// helper method for std::sort of outputs
bool sortOutputs(const CTxOut &a, const CTxOut &b) {
    if (a.nValue < b.nValue)
        return true;
    if (a.nValue > b.nValue)
        return false;
    return a.scriptPubKey < b.scriptPubKey;
}

// if anonimize is true, this will apply
// BIP69 sorting that aims to anonimize the transaction.
// any sorting of inputs is also applied to the signInfo vector.
CMutableTransaction sortTransaction(const CMutableTransaction &tx, std::vector<TransactionBuilderPrivate::SignInfo> &signInfo)
{
    assert(tx.vin.size() == signInfo.size());

    // First check if the transaction is eligable for sorting.
    for (size_t i = 0; i < tx.vin.size(); ++i) {
        // the 0x40 is the replay protection between BTC and BCH
        // then we also allow SignAllOuputs  and SignNoOutputs
        constexpr uint8_t BadMask = 0xBC; // inverted of the allowed

        const auto &si = signInfo.at(i);
        uint8_t hashType = si.hashType;
        if (si.prevOutScript.empty()) {
            // we won't sign this, use the hashtype already on the Tx.
            const auto &script = tx.vin.at(i).scriptSig;
            if (script.empty()) // no signature, no problem
                continue;

            auto scriptIter = script.begin();
            opcodetype type;
            std::vector<unsigned char> p2pkhSignature;
            script.GetOp(scriptIter, type, p2pkhSignature);
            // for the type we only allow PUSHDATA2 or a push by-size (i.e a byte less than PUSHDATA1)
            if (type == OP_PUSHDATA1
                    || type <= 0
                    || type > OP_PUSHDATA2
                    || p2pkhSignature.size() < 65) // custom unlock script.
                return tx; // we wont't sort that, it will likely break the signature.

            hashType = p2pkhSignature.back();
        }
        if ((hashType & BadMask) != 0) // can't sort this one.
            return tx;
    }

    // ok, lets make a sorted copy!
    const std::vector<TransactionBuilderPrivate::SignInfo> origSignInfo(signInfo);

    std::vector<int> inputs;
    inputs.reserve(tx.vin.size());
    for (size_t i = 0; i < tx.vin.size(); ++i) inputs.push_back(static_cast<int>(i));
    InputSorter inputSorter(tx);
    std::sort(inputs.begin(), inputs.end(), inputSorter);

    CMutableTransaction out;
    out.vin.resize(inputs.size());
    for (size_t i = 0; i < inputs.size(); ++i) {
        const size_t sorted = static_cast<size_t>(inputs[i]);
        out.vin[i] = tx.vin.at(sorted);
        signInfo[i] = origSignInfo.at(sorted);
    }

    // sort outputs
    out.vout = tx.vout;
    std::sort(out.vout.begin(), out.vout.end(), sortOutputs);
    return out;
}

}



TransactionBuilder::TransactionBuilder()
    : d(new TransactionBuilderPrivate())
{
}

TransactionBuilder::TransactionBuilder(const Tx &existingTx)
    : d(new TransactionBuilderPrivate())
{
    d->transaction = CTransaction(existingTx.createOldTransaction());
    d->signInfo.resize(d->transaction.vin.size());
}

TransactionBuilder::TransactionBuilder(const CTransaction &existingTx)
    : d(new TransactionBuilderPrivate())
{
    d->transaction = CTransaction(existingTx);
    d->signInfo.resize(d->transaction.vin.size());
}

TransactionBuilder::~TransactionBuilder()
{
    delete d;
}

int TransactionBuilder::appendInput(const uint256 &txid, int outputIndex)
{
    const size_t pos = d->transaction.vin.size();
    if (pos > 1000) // kind of random large number
        throw std::runtime_error("Too many inputs");
    d->transaction.vin.resize(pos + 1);
    d->signInfo.resize(pos + 1);
    CTxIn &in = d->transaction.vin[pos];
    in.prevout.hash = txid;
    in.prevout.n = outputIndex;
    switch (d->defaultLocking) {
    case TransactionBuilder::LockMiningOnTime:
    case TransactionBuilder::LockMiningOnBlock:
        in.nSequence = in.SEQUENCE_LOCKTIME_DISABLE_FLAG;
        break;
    default: // default of the instance is fine
        break;
    }
    d->curInput = static_cast<int>(pos);
    return d->curInput;
}

int TransactionBuilder::selectInput(int index)
{
    assert(index >= 0);
    if (index < 0) throw std::runtime_error("Index is a natural number");
    d->curInput = std::min(static_cast<int>(d->transaction.vin.size()) -1, index);
    return d->curInput;
}

int TransactionBuilder::outputCount() const
{
    return static_cast<int>(d->transaction.vout.size());
}

int TransactionBuilder::inputCount() const
{
    return static_cast<int>(d->transaction.vin.size());
}

void TransactionBuilder::pushInputSignature(const PrivateKey &privKey, const CScript &prevOutScript, int64_t amount, SignatureType type, SignInputs inputs, SignOutputs outputs)
{
    d->checkCurInput();
    TransactionBuilderPrivate::SignInfo &si = d->signInfo[d->curInput];
    si.hashType = inputs == SignOnlyThisInput ? 0xC0 : 0x40;
    switch (outputs) {
    case SignAllOuputs: si.hashType += 1; break;
    case SignNoOutputs: si.hashType += 2; break;
    case SignSingleOutput: si.hashType += 3; break;
    }

    si.privKey = privKey;
    si.prevOutScript = prevOutScript;
    si.amount = amount;
    si.signatureType = type;
}

void TransactionBuilder::deleteInput(int index)
{
    assert(index >= 0);
    assert(size_t(index) < d->transaction.vin.size());
    assert(size_t(index) < d->signInfo.size());
    auto iter = d->transaction.vin.begin();
    iter += index;
    d->transaction.vin.erase(iter);

    auto iter2 = d->signInfo.begin();
    iter2 += index;
    d->signInfo.erase(iter2);

    selectInput(index);
}

int TransactionBuilder::appendOutput(int64_t amount)
{
    const size_t pos = d->transaction.vout.size();
    if (pos > 1000) // kind of random large number
        throw std::runtime_error("Too many outputs");
    d->transaction.vout.resize(pos + 1);
    CTxOut &out = d->transaction.vout[pos];
    out.nValue = amount;

    d->curOutput = static_cast<int>(pos);
    return d->curOutput;
}

int TransactionBuilder::selectOutput(int index)
{
    assert(index >= 0);
    if (index < 0) throw std::runtime_error("Index is a natural number");
    d->curOutput = std::min(static_cast<int>(d->transaction.vout.size()) -1, index);
    return d->curOutput;
}

void TransactionBuilder::setOutputValue(int64_t value)
{
    assert(value >= 0);
    assert(d->curOutput >= 0);
    assert(int(d->transaction.vout.size()) > d->curOutput);
    d->transaction.vout[d->curOutput].nValue = value;
}

void TransactionBuilder::pushOutputPay2Address(const KeyId &address)
{
    d->checkCurOutput();
    CScript outScript;
    outScript << OP_DUP << OP_HASH160;
    std::vector<unsigned char> data(address.begin(), address.end());
    outScript << data;
    outScript << OP_EQUALVERIFY << OP_CHECKSIG;
    pushOutputScript(outScript);
}

void TransactionBuilder::pushOutputScript(const CScript &script)
{
    assert(d->curOutput >= 0);
    assert(static_cast<size_t>(d->curOutput) < d->transaction.vout.size());
    d->transaction.vout[static_cast<size_t>(d->curOutput)].scriptPubKey = script;
}

void TransactionBuilder::pushOutputPay2Script(const CScriptID &script)
{
    d->checkCurOutput();
    CScript outScript;
    outScript << OP_HASH160;
    std::vector<unsigned char> data(script.begin(), script.end());
    outScript << data;
    outScript << OP_EQUAL;
    d->transaction.vout[d->curOutput].scriptPubKey = outScript;
}

void TransactionBuilder::pushOutputNullData(const std::vector<uint8_t> &data)
{
    assert(d->curOutput >= 0);
    assert(static_cast<size_t>(d->curOutput) < d->transaction.vout.size());
    CScript script;
    script << OP_RETURN << data;
    d->transaction.vout[static_cast<size_t>(d->curOutput)].scriptPubKey = script;
}

void TransactionBuilder::deleteOutput(int index)
{
    assert(index >= 0);
    assert(size_t(index) < d->transaction.vout.size());
    auto iter = d->transaction.vout.begin();
    iter+=index;
    d->transaction.vout.erase(iter);
    selectOutput(index);
}

Tx TransactionBuilder::createTransaction(const std::shared_ptr<Streaming::BufferPool> &pool)
{
    CMutableTransaction *tx = &d->transaction;
    CMutableTransaction dummy;
    std::vector<TransactionBuilderPrivate::SignInfo> signInfo(d->signInfo);
    if (d->anonimize) {
        // sort the transaction prior to signing inputs.
        dummy = sortTransaction(d->transaction, signInfo);
        assert(signInfo.size() == d->signInfo.size());
        tx = &dummy;
    }

    // sign all inputs we can.
    assert(tx->vin.size() == signInfo.size());
    for (size_t i = 0; i < tx->vin.size(); ++i) {
        const TransactionBuilderPrivate::SignInfo &si = signInfo[i];
        if (si.prevOutScript.empty())
            continue;

        uint256 hashPrevouts;
        if (!(si.hashType & SignOnlyThisInput)) {
            CHashWriter ss(SER_GETHASH, 0);
            for (size_t n = 0; n < tx->vin.size(); ++n) {
                ss << tx->vin[n].prevout;
            }
            hashPrevouts = ss.finalizeHash();
        }
        uint256 hashSequence;
        if (!(si.hashType & SignOnlyThisInput) && (si.hashType & 0x1f) != SignSingleOutput
                && (si.hashType & 0x1f) != SignNoOutputs) {
            CHashWriter ss(SER_GETHASH, 0);
            for (size_t n = 0; n < tx->vin.size(); ++n) {
                ss << tx->vin[n].nSequence;
            }
            hashSequence = ss.finalizeHash();
        }
        uint256 hashOutputs;
        if ((si.hashType & 0x1f) != SignSingleOutput && (si.hashType & 0x1f) != SignNoOutputs) {
            CHashWriter ss(SER_GETHASH, 0);
            for (size_t n = 0; n < tx->vout.size(); ++n) {
                ss << tx->vout[n];
            }
            hashOutputs = ss.finalizeHash();
        } else if ((si.hashType & 0x1f) == SignSingleOutput && i < tx->vout.size()) {
            CHashWriter ss(SER_GETHASH, 0);
            ss << tx->vout[i];
            hashOutputs = ss.finalizeHash();
        }

        // use FORKID based creation of the hash we will sign.
        CHashWriter ss(SER_GETHASH, 0);
        ss << tx->nVersion << hashPrevouts << hashSequence;
        ss << tx->vin[i].prevout;
        ss << static_cast<const CScriptBase &>(si.prevOutScript);
        ss << si.amount << tx->vin[i].nSequence << hashOutputs;
        ss << tx->nLockTime << (int) si.hashType;
        const uint256 hash = ss.finalizeHash();

        // the rest assumes P2PKH for now.
        std::vector<unsigned char> vchSig;
        if (si.signatureType == ECDSA)
            si.privKey.signECDSA(hash, vchSig);
        else
            si.privKey.signSchnorr(hash, vchSig);
        vchSig.push_back((uint8_t) si.hashType);

        tx->vin[i].scriptSig = CScript();
        tx->vin[i].scriptSig << vchSig;
        tx->vin[i].scriptSig << ToByteVector(si.privKey.getPubKey());
    }
    return Tx::fromOldTransaction(*tx, pool);
}

bool TransactionBuilder::anonimize() const
{
    return d->anonimize;
}

void TransactionBuilder::setAnonimize(bool on)
{
    d->anonimize = on;
}