// Copyright (c) 2025 The Limenka developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#ifndef LIMENKA_CONSENSUS_CT_H
#define LIMENKA_CONSENSUS_CT_H

#include <hash.h>
#include <primitives/transaction.h>
#include <span.h>
#include <uint256.h>

#include <cstddef>
#include <cstdint>
#include <optional>
#include <string>

// Confidential transaction (P2BPCT) kernel output magic: OP_RETURN <"BK">.
static constexpr uint8_t CT_KERNEL_MAGIC_BYTE0 = 0x42; // 'B'
static constexpr uint8_t CT_KERNEL_MAGIC_BYTE1 = 0x4B; // 'K'

// Marker for "no kernel output found in this transaction".
static constexpr int NO_CT_KERNEL_OUTPUT = -1;

// Kernel signature size (BIP340 Schnorr) and fee field size.
// The fee is a 128-bit value in the same units as the committed amounts.
// Stealth fields: a 1-byte marker, a 33-byte ephemeral point E, and two
// 32-byte encrypted fields (amount to the view key, blinding to the spend
// key).
static constexpr size_t CT_KERNEL_SIG_SIZE = 64;
static constexpr size_t CT_FEE_SIZE = 16;
static constexpr size_t CT_STEALTH_EPHEM_SIZE = 33;
static constexpr uint8_t CT_STEALTH_MARKER = 0x53; // 'S'
static constexpr size_t CT_STEALTH_ENC_SIZE = 32;

// Kernel output format:
//   OP_RETURN <"BK"> <fee:16 little-endian> <sig:64>
//             [<"S"> <E:33> <enc_amount:32> <enc_blind:32>]
//
// Nothing else may follow: the kernel format is closed.  The balance
// equation is a plain Pedersen point sum - there are NO offset terms.
// (A previous design carried "blinding offsets" as arbitrary points;
// that allowed an attacker to include b*H points for chosen b and mint
// arbitrary value, since consensus cannot distinguish a G-only point
// from an H-containing one.  Offsets are removed entirely.)
//
// Stealth payments (receiver publishes a static (view, spend) address):
//   The sender chooses the blinding r itself (so it can sign the kernel
//   excess that covers the output - no receiver interaction, no offset
//   points), then encrypts the amount to the view key and the blinding
//   to the spend key:
//   E         = t*G (ephemeral, shared secret basis)
//   enc_amount= v XOR H("CTStealthAmount" || t*V_view)
//   enc_blind = r XOR H("CTStealthBlind" || t*V_spend)
//   The view key alone can scan (decrypt v); only the spend key can
//   recover the blinding.  The commitment C = r*G + v*H is an ordinary
//   confidential output bound by the balance equation.
struct CTKernelData {
    CAmount fee{0};
    std::vector<uint8_t> sig; // 64 bytes
    bool has_stealth{false};
    std::vector<uint8_t> E;          // 33 bytes (empty if !has_stealth)
    std::vector<uint8_t> enc_amount; // 32 bytes (empty if !has_stealth)
    std::vector<uint8_t> enc_blind;  // 32 bytes (empty if !has_stealth)
};

/** True if the script is the CT kernel output prefix: OP_RETURN <"BK"> ... */
inline bool IsCTKernelScript(const CScript& spk)
{
    if (spk.size() < 4) return false;
    if (spk[0] != OP_RETURN) return false;
    if (spk[1] != 0x02) return false; // push 2 bytes
    return spk[2] == CT_KERNEL_MAGIC_BYTE0 && spk[3] == CT_KERNEL_MAGIC_BYTE1;
}

/** Find the index of the confidential transaction kernel output.
 *  Returns NO_CT_KERNEL_OUTPUT if not found (takes the last match).
 */
inline int GetCTKernelOutputIndex(const CTransaction& tx)
{
    int pos = NO_CT_KERNEL_OUTPUT;
    for (size_t o = 0; o < tx.vout.size(); ++o) {
        const CScript& spk = tx.vout[o].scriptPubKey;
        CScript::const_iterator pc = spk.begin();
        opcodetype opcode;
        std::vector<uint8_t> data;
        if (!spk.GetOp(pc, opcode, data) || opcode != OP_RETURN) continue;
        if (!spk.GetOp(pc, opcode, data)) continue;
        if (data.size() == 2 && data[0] == CT_KERNEL_MAGIC_BYTE0 &&
            data[1] == CT_KERNEL_MAGIC_BYTE1) {
            pos = static_cast<int>(o);
        }
    }
    return pos;
}

/** Parse the fee and kernel signature from a kernel output. */
inline std::optional<CTKernelData> ParseCTKernelOutput(const CTxOut& vout)
{
    const CScript& spk = vout.scriptPubKey;
    CScript::const_iterator pc = spk.begin();
    opcodetype opcode;
    std::vector<uint8_t> magic;
    std::vector<uint8_t> fee;
    std::vector<uint8_t> sig;
    if (!spk.GetOp(pc, opcode, magic) || opcode != OP_RETURN) return std::nullopt;
    if (!spk.GetOp(pc, opcode, magic)) return std::nullopt;
    if (magic.size() != 2 || magic[0] != CT_KERNEL_MAGIC_BYTE0 ||
        magic[1] != CT_KERNEL_MAGIC_BYTE1) return std::nullopt;
    if (!spk.GetOp(pc, opcode, fee) || fee.size() != CT_FEE_SIZE) return std::nullopt;
    if (!spk.GetOp(pc, opcode, sig) || sig.size() != CT_KERNEL_SIG_SIZE) return std::nullopt;

    // Optional stealth fields: <"S"> <E:33> <enc_amount:32> <enc_blind:32>.
    bool has_stealth = false;
    std::vector<uint8_t> E, enc_amount, enc_blind;
    {
        std::vector<uint8_t> marker;
        if (spk.GetOp(pc, opcode, marker) && marker.size() == 1 &&
            marker[0] == CT_STEALTH_MARKER) {
            if (!spk.GetOp(pc, opcode, E) || E.size() != CT_STEALTH_EPHEM_SIZE)
                return std::nullopt;
            if (!spk.GetOp(pc, opcode, enc_amount) || enc_amount.size() != CT_STEALTH_ENC_SIZE)
                return std::nullopt;
            if (!spk.GetOp(pc, opcode, enc_blind) || enc_blind.size() != CT_STEALTH_ENC_SIZE)
                return std::nullopt;
            has_stealth = true;
        } else if (!marker.empty()) {
            // The kernel format is closed: no fields past the stealth block.
            return std::nullopt;
        }
    }

    // The kernel format is closed: no trailing fields.
    {
        std::vector<uint8_t> extra;
        if (spk.GetOp(pc, opcode, extra)) return std::nullopt;
    }

    CAmount fee_val = 0;
    for (size_t i = 0; i < CT_FEE_SIZE; i++) {
        // Shift through unsigned so byte 15 cannot shift into the sign bit
        // (signed left shift into the sign bit is UB).
        fee_val |= static_cast<CAmount>(static_cast<__uint128_t>(fee[i]) << (8 * i));
    }

    CTKernelData out;
    out.fee = fee_val;
    out.sig = std::move(sig);
    out.has_stealth = has_stealth;
    out.E = std::move(E);
    out.enc_amount = std::move(enc_amount);
    out.enc_blind = std::move(enc_blind);
    return out;
}

/** Compute the 32-byte kernel message the excess signature signs.
 *  Binds the fee, input prevouts, output scriptPubKeys (excluding the
 *  kernel output, whose scriptPubKey contains the signature itself), and
 *  the stealth fields (so they cannot be substituted).
 */
inline uint256 ComputeCTKernelMessage(const CTransaction& tx, int kernel_index,
                                      const CTKernelData& kernel,
                                      const std::vector<CAmount>& transparent_input_values = {})
{
    HashWriter hw{};
    hw << std::string("ct-kernel");
    // 16-byte little-endian fee (the __int128 stream serializer is 8-byte,
    // so write the two halves explicitly).
    hw << static_cast<uint64_t>(kernel.fee)
       << static_cast<uint64_t>(static_cast<__uint128_t>(kernel.fee) >> 64);
    for (const auto& txin : tx.vin) {
        hw << txin.prevout;
    }
    // Mint inputs: bind their visible values (attosats) so the balance
    // cannot be re-targeted after signing.
    hw << static_cast<uint64_t>(transparent_input_values.size());
    for (const CAmount v : transparent_input_values) {
        hw << static_cast<uint64_t>(v)
           << static_cast<uint64_t>(static_cast<__uint128_t>(v) >> 64);
    }
    for (size_t i = 0; i < tx.vout.size(); i++) {
        if (static_cast<int>(i) == kernel_index) continue;
        hw << tx.vout[i].scriptPubKey;
    }
    hw << static_cast<uint8_t>(kernel.has_stealth ? 1 : 0);
    if (kernel.has_stealth) {
        hw << kernel.E;
        hw << kernel.enc_amount;
        hw << kernel.enc_blind;
    }
    return hw.GetHash();
}

#endif // LIMENKA_CONSENSUS_CT_H