/*
 * This file is part of the Flowee project
 * Copyright (C) 2019-2024 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>
#include <streaming/StreamingUtils.h>

static constexpr size_t MaxCashTokenCommitmentSize = 128;

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 value = 0;
        PrivateKey privKey;
        CScript prevOutScript;
        TransactionBuilder::SignatureType signatureType = TransactionBuilder::Schnorr;
    };
    std::vector<SignInfo> signInfo;

    struct CashToken {
        int outputIndex = -1;
        uint8_t bitfield = 0;
        uint256 category;
        uint64_t amount = 0;
        std::vector<uint8_t> commitment;
    };
    std::vector<CashToken> cashTokens;

    bool getToken(int outputIndex, CashToken **token, bool autocreate = true) {
        for (size_t i = 0; i < cashTokens.size(); ++i) {
            if (cashTokens.at(i).outputIndex == outputIndex) {
                *token = &cashTokens[i];
                return true;
            }
        }
        if (!autocreate)
            return false;
        cashTokens.resize(cashTokens.size() + 1);
        *token = &cashTokens.back();
        assert(*token);
        (*token)->outputIndex = outputIndex;
        return true;
    }

    bool anonimize = false;
    int feePerKByte = 1000;
    int outputToTakeFeeFrom = -1;
};

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 struct to sort index-of-outputs
struct OutputSorter {
    const CMutableTransaction &tx;
    explicit OutputSorter(const CMutableTransaction &transaction) : tx(transaction) {}
    bool operator()(int a, int b) const {
        assert(a >= 0);
        assert(b >= 0);
        assert((int)tx.vout.size() > a);
        assert((int)tx.vout.size() > b);
        const CTxOut &oa = tx.vout.at(a);
        const CTxOut &ob = tx.vout.at(b);

        if (oa.nValue < ob.nValue)
            return true;
        if (oa.nValue > ob.nValue)
            return false;
        return oa.scriptPubKey < ob.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, std::vector<int> &outSortMap)
{
    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);

    // first inputs. We sort a vector of ints based on the content (using InputSorter)
    // and then copy based on that.
    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);
    }

    // Same again, but now for outputs.
    outSortMap.resize(0);
    outSortMap.reserve(tx.vout.size());
    for (size_t i = 0; i < tx.vout.size(); ++i) outSortMap.push_back(static_cast<int>(i));
    OutputSorter outputSorter(tx);
    std::sort(outSortMap.begin(), outSortMap.end(), outputSorter);

    out.vout.resize(outSortMap.size());
    for (size_t i = 0; i < outSortMap.size(); ++i) {
        const size_t sorted = static_cast<size_t>(outSortMap[i]);
        out.vout[i] = tx.vout.at(sorted);
    }
    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 = existingTx;
    d->signInfo.resize(d->transaction.vin.size());
}

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

void TransactionBuilder::setFeeTarget(int satsPerKByte)
{
    d->feePerKByte = satsPerKByte;
}

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 value, 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.value = value;
    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 value)
{
    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 = value;

    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::setOutputFeeSource(int outputIndex)
{
    if (outputIndex == -1)
       outputIndex = d->curOutput;
    d->outputToTakeFeeFrom = outputIndex;
}

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 std::vector<uint8_t> &p2SHash)
{
    d->checkCurOutput();
    CScript outScript;
    if (p2SHash.size() == 20)
        outScript << OP_HASH160;
    else if (p2SHash.size() == 32)
        outScript << OP_HASH256;
    outScript << p2SHash;
    outScript << OP_EQUAL;
    pushOutputScript(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);
    if (size_t(index) >= d->transaction.vout.size())
        throw std::runtime_error("Out of bounds");
    auto iter = d->transaction.vout.begin();
    iter += index;
    d->transaction.vout.erase(iter);

    // adjust indexes on cashTokens
    for (auto iter = d->cashTokens.begin(); iter != d->cashTokens.end(); ++iter) {
        if (iter->outputIndex == index)
            d->cashTokens.erase(iter);
        else if (iter->outputIndex > index)
            --iter->outputIndex;
    }

    if (d->outputToTakeFeeFrom >= index)
        --d->outputToTakeFeeFrom;

    selectOutput(index);
}

void TransactionBuilder::startOutputToken(const uint256 &category, uint8_t bitfield)
{
    assert(d->curOutput >= 0);
    TransactionBuilderPrivate::CashToken *ct = nullptr;
    d->getToken(d->curOutput, &ct);
    assert(ct);
    ct->category = category;
    ct->bitfield = bitfield;
}

void TransactionBuilder::pushNftCommitment(const std::vector<uint8_t> &data)
{
    assert(d->curOutput >= 0);
    if (data.size() > MaxCashTokenCommitmentSize)
        throw std::runtime_error("Illegal commitment size");
    TransactionBuilderPrivate::CashToken *ct = nullptr;
    if (!d->getToken(d->curOutput, &ct, false))
        throw std::runtime_error("Builder: token not started");
    assert(ct);
    if ((ct->bitfield & Tx::HasCommitment) == 0)
        throw std::runtime_error("push: Token has no commitment");
    ct->commitment = data;
}

void TransactionBuilder::pushOutputFtAmount(uint64_t amount)
{
    assert(d->curOutput >= 0);
    TransactionBuilderPrivate::CashToken *ct = nullptr;
    if (!d->getToken(d->curOutput, &ct, false))
        throw std::runtime_error("Builder: token not started");
    assert(ct);
    if ((ct->bitfield & Tx::HasFtAmount) == 0)
        throw std::runtime_error("push: Token has no FT");
    ct->amount = amount;
}

void TransactionBuilder::clearOutputToken()
{
    for (auto iter = d->cashTokens.begin(); iter != d->cashTokens.end(); ++iter) {
        if (iter->outputIndex == d->curOutput) {
            d->cashTokens.erase(iter);
            return;
        }
    }
}

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

    // add cashtokens to outputs
    for (const auto &ct : d->cashTokens) {
        assert(ct.outputIndex >= 0);
        size_t outputIndex = ct.outputIndex;
        if (d->anonimize) {
            assert(outputIndex < outSortMap.size());
            outputIndex = outSortMap.at(outputIndex);
        }

        assert(tx.vout.size() > outputIndex);
        auto base = tx.vout[outputIndex].scriptPubKey;
        std::vector<uint8_t> output;
        output.reserve(32 + MaxCashTokenCommitmentSize + 10 + base.size());
        output.resize(34);
        output[0] = 0xef; // magic indicator
        memcpy(&output[1], ct.category.begin(), 32);
        output[33] = ct.bitfield;
        if (ct.bitfield & Tx::HasCommitment) {
            assert(ct.commitment.size() <= MaxCashTokenCommitmentSize);
            if (ct.commitment.size() == 0)
                throw std::runtime_error("Token commitment missing");
            const size_t i = output.size();
            output.resize(i + ct.commitment.size() + 1);
            output[i] = (uint8_t) ct.commitment.size();
            memcpy(&output[i + 1], &ct.commitment[0], ct.commitment.size());
        }
        if (ct.bitfield & Tx::HasFtAmount) {
            if (ct.amount == 0 || ct.amount > 0xffffffffffffff7f)
                throw std::runtime_error("Token Ft-amount wrong");

            const size_t i = output.size();
            output.resize(output.size() + 9);
            auto bytes = Streaming::writeCompactSize(reinterpret_cast<char*>(&output[i]), ct.amount);
            output.resize(i + bytes);
        }
        const size_t i = output.size();
        output.resize(output.size() + base.size());
        memcpy(&output[i], &base[0], base.size());

        tx.vout[outputIndex].scriptPubKey = CScript(output.begin(), output.end());
    }

    /* sign and do fee
     *
     * The fee is set based on the size of the transaction. The main way to change the size is with a different
     * signature. Signatures change based on the fee if that comes from an existing output. So, this requires
     * iteration and at least 2 passes to make work.
     */
    while (true) {
        // Sign all inputs we can.
        bool haveAllInputAmounts = true;
        int64_t totalInputs = 0;
        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()) {
                haveAllInputAmounts = false;
                continue;
            }
            totalInputs += si.value;
            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.value << 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());
        }

        // now lets see if fee adjustment is needed (and wanted)
        if (!(d->feePerKByte > 0 && d->outputToTakeFeeFrom >= 0 && d->outputToTakeFeeFrom < (int)tx.vout.size() && haveAllInputAmounts)) {
            // ok, no, we won't be updating fee.
            break;
        }

        // We'll adjust the output amount of one output in order to reach the expected number of sats per kbyte.
        int outputToUse = d->outputToTakeFeeFrom;
        if (d->anonimize) {
            assert(outputToUse < (int) outSortMap.size());
            outputToUse = outSortMap.at(outputToUse);
        }

        int assignedFee = totalInputs;
        for (const auto &out : tx.vout) {
            assignedFee -= out.nValue;
        }
        // a positive diff means we underpaid fee
        const int diff = static_cast<int64_t>(tx.GetSerializeSize(1, 1)) * d->feePerKByte / 1000 - assignedFee;
        if (diff <= 0)
            break; // done

        tx.vout[outputToUse].nValue -= diff;
        // We're inside of a while-true, so lets re-calc the signatures!
    }

    return Tx::fromOldTransaction(tx, pool);
}

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

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

void TransactionBuilder::setTransactionVersion(int txVersion)
{
    d->transaction.nVersion = txVersion;
}

int TransactionBuilder::transactionVersion() const
{
    return d->transaction.nVersion;
}