// 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_WALLET_CT_H
#define LIMENKA_WALLET_CT_H

#include <consensus/ct.h>
#include <psbt.h>
#include <crypto/bulletproofs.h>
#include <key.h>
#include <pubkey.h>
#include <primitives/transaction.h>
#include <random.h>
#include <uint256.h>
#include <util/result.h>

#include <cstdint>
#include <utility>
#include <vector>

namespace wallet {

class CWallet;

// One lambda (whole unit) in attosats: 10^26 = satoshis (10^8) * attosats
// per satoshi (10^18).
static constexpr CAmount LAMBDA_SCALE = CAmount{100000000} * ATTOSATS_PER_SATOSHI;

// Format an attosat amount as lambda with up to 26 decimals (trailing
// zeros trimmed).  Parsing accepts fixed-point strings of the same shape.
std::string AttosatsToString(CAmount attosats);
bool ParseAttosatsString(const std::string& s, CAmount& out);

// HD derivation paths for the wallet's stealth keypair (dedicated purpose).
// Both hardened so the master key suffices to recover them.
static constexpr uint32_t STEALTH_VIEW_PATH = 0x800001F4;  // m/500'
static constexpr uint32_t STEALTH_SPEND_PATH = 0x800001F5; // m/501'

// A stealth CT address: the receiver publishes (view, spend) once and the
// sender pays to it non-interactively.  The view key recovers the blinding
// via ECDH with the kernel ephemeral; the spend key separates scanning
// capability (watchtowers, lightning monitors) from spending capability.
struct StealthCTAddress {
    CPubKey view;   //!< V_scan (33 bytes compressed)
    CPubKey spend;  //!< V_spend (33 bytes compressed)

    SERIALIZE_METHODS(StealthCTAddress, obj) { READWRITE(obj.view, obj.spend); }
};

// Sender-side result of deriving a stealth output.
struct StealthPayment {
    BPCommitment commitment;        //!< C = blind*G + v*H (33 bytes)
    BPScalar blinding;              //!< sender-chosen r (sender signs the excess)
    std::vector<uint8_t> E;         //!< ephemeral t*G (33 bytes)
    std::vector<uint8_t> enc_amount;//!< v XOR H(shared_v) (32 bytes)
    std::vector<uint8_t> enc_blind; //!< r XOR H(shared_s) (32 bytes)
    CAmount amount{0};              //!< attosats
};

// Everything the wallet must retain to spend a confidential output it
// created.  The bulletproof rides in the spend witness; the blinding is
// needed to compute the kernel excess of the spending transaction.
struct CTReceipt {
    uint64_t amount_lo{0}; //!< attosats, little-endian halves (16 bytes total)
    uint64_t amount_hi{0};
    BPScalar blinding;      //!< 32 bytes, fresh per output
    BPScalar seed;          //!< 32 bytes, fresh per output
    uint32_t vout_index{0}; //!< position of the output in its creating tx
    Bulletproof proof;      //!< 754 bytes serialized

    SERIALIZE_METHODS(CTReceipt, obj) {
        READWRITE(obj.amount_lo, obj.amount_hi, obj.blinding, obj.seed,
                  obj.vout_index,
                  obj.proof.A, obj.proof.S, obj.proof.T1, obj.proof.T2,
                  obj.proof.t_hat, obj.proof.taux, obj.proof.mu, obj.proof.a,
                  obj.proof.b);
        for (size_t i = 0; i < BP_ROUNDS; ++i) {
            READWRITE(obj.proof.L[i], obj.proof.R[i]);
        }
    }

    CAmount Amount() const
    {
        return static_cast<CAmount>(amount_lo) |
               (static_cast<CAmount>(amount_hi) << 64);
    }
    void SetAmount(CAmount v)
    {
        amount_lo = static_cast<uint64_t>(v);
        amount_hi = static_cast<uint64_t>(static_cast<__uint128_t>(v) >> 64);
    }
};

// Result of building a confidential transaction.
struct CTCreationResult {
    CMutableTransaction tx;
    std::vector<CTReceipt> new_receipts; // one per output, same order
    int kernel_index{-1};
    // Present when this is a stealth payment: the sender-chosen blinding
    // and encrypted fields (the blinding must fold into the PSBT excess).
    std::optional<StealthPayment> stealth_payment;
};

// The witness v4 output script for a committed amount: OP_4 <commitment>.
CScript GetConfidentialScript(const BPCommitment& commitment);

// Create a confidential output committing amount_attosats.  Derives a
// fresh blinding and seed from the CSPRNG, proves the range [0, 2^128),
// and fills the receipt the wallet must keep.
bool CreateConfidentialOutput(CAmount amount_attosats, FastRandomContext& rng,
                              BPCommitment& commitment, CTReceipt& receipt);

// excess = sum(in_blindings) - sum(out_blindings) mod n.  Returns false if
// the excess is zero (the kernel key would be the identity - re-blind an
// output and try again).
bool ComputeCTExcess(const std::vector<BPScalar>& in_blindings,
                     const std::vector<BPScalar>& out_blindings, BPScalar& excess);

// Build the kernel output script:
//   OP_RETURN <"BK"> <fee:16 LE> <sig:64> [<"S"> <E> <enc_amount> <enc_blind>]
// The message binds the fee, input prevouts, output scripts, and the
// stealth fields (substitution protection).
bool BuildCTKernelOutput(const BPScalar& excess, CAmount fee, const CTransaction& tx,
                         int kernel_index,
                         const std::vector<uint8_t>& E,
                         const std::vector<uint8_t>& enc_amount,
                         const std::vector<uint8_t>& enc_blind,
                         CScript& kernel_script,
                         const std::vector<CAmount>& transparent_input_values = {});

// Derive the stealth address from the receiver's view and spend secrets.
StealthCTAddress CreateStealthCTAddress(const CKey& view_secret, const CKey& spend_secret);

// Sender side: derive a stealth output for the receiver's static address.
// The sender chooses the blinding r itself, so its excess signature covers
// the output without any receiver interaction (no kernel offset points).
// C = r*G + v*H; E = t*G; enc_amount = v XOR H(shared_v) with
// shared_v = t*V_view; enc_blind = r XOR H(shared_s) with
// shared_s = t*V_spend.  The view key can scan (decrypt v) but only the
// spend key can recover the blinding.
bool CreateStealthOutput(CAmount amount_attosats, const StealthCTAddress& address,
                         FastRandomContext& rng, StealthPayment& out);

// Receiver side: recover (amount, blinding) from a kernel's stealth
// fields.  Decrypts v with the view key and r with the spend key, then
// verifies the commitment C == r*G + v*H.  Returns nullopt if the output
// does not belong to this receiver.
std::optional<std::pair<CAmount, BPScalar>> RecoverStealthOutput(
    const CKey& view_secret, const CKey& spend_secret,
    const std::vector<uint8_t>& E, const std::vector<uint8_t>& enc_amount,
    const std::vector<uint8_t>& enc_blind, const BPCommitment& commitment);

// Build a confidential transaction paying a stealth address.  Inputs are
// CT outputs the wallet holds receipts for; the kernel carries the
// ephemeral and the encrypted amount/blinding fields.
util::Result<CTCreationResult> CreateStealthTransaction(
    const std::vector<std::pair<COutPoint, CTReceipt>>& ct_inputs,
    const StealthCTAddress& address, CAmount amount, CAmount fee, FastRandomContext& rng,
    int change_outputs = 1);

// Split a total attosat value into n random, widely-varying parts (for
// privacy: temporal rejoining is harder when the sizes differ).  The
// last part carries the remainder.  n is clamped to [1, 16].
std::vector<CAmount> SplitAttosats(CAmount total, int n, FastRandomContext& rng);

// PSBT global-proprietary payload for confidential transactions: the
// kernel metadata (signature empty until finalize) plus the blindings and
// proofs needed to sign the kernel.  The blindings are secret - a CT PSBT
// must be treated as sensitive, like a PSBT carrying private keys.
struct CTPSBTData {
    CTKernelData kernel;                  //!< fee/stealth fields; sig empty
    std::vector<CTReceipt> in_receipts;   //!< one per CT input (blinding+proof)
    std::vector<CTReceipt> out_receipts;  //!< one per CT output (blinding+proof)
    //! Blindings of outputs the wallet built but does not own (stealth
    //! payments to lm2): they fold into the kernel excess but carry no
    //! receipt (the receiver generates its own).
    std::vector<BPScalar> out_blindings_extra;

    template <typename Stream>
    void Serialize(Stream& s) const
    {
        const uint64_t fee_lo = static_cast<uint64_t>(kernel.fee);
        const uint64_t fee_hi = static_cast<uint64_t>(static_cast<__uint128_t>(kernel.fee) >> 64);
        ::Serialize(s, fee_lo);
        ::Serialize(s, fee_hi);
        ::Serialize(s, kernel.has_stealth);
        ::Serialize(s, kernel.E);
        ::Serialize(s, kernel.enc_amount);
        ::Serialize(s, kernel.enc_blind);
        ::Serialize(s, kernel.sig);
        ::Serialize(s, in_receipts);
        ::Serialize(s, out_receipts);
        ::Serialize(s, out_blindings_extra);
    }

    template <typename Stream>
    void Unserialize(Stream& s)
    {
        uint64_t fee_lo, fee_hi;
        ::Unserialize(s, fee_lo);
        ::Unserialize(s, fee_hi);
        kernel.fee = static_cast<CAmount>(fee_lo) | (static_cast<CAmount>(fee_hi) << 64);
        ::Unserialize(s, kernel.has_stealth);
        ::Unserialize(s, kernel.E);
        ::Unserialize(s, kernel.enc_amount);
        ::Unserialize(s, kernel.enc_blind);
        ::Unserialize(s, kernel.sig);
        ::Unserialize(s, in_receipts);
        ::Unserialize(s, out_receipts);
        ::Unserialize(s, out_blindings_extra);
    }
};

// Store/extract the CT payload in a PSBT global proprietary field.
bool SetCTPSBTData(PartiallySignedTransaction& psbt, const CTPSBTData& data);
bool GetCTPSBTData(const PartiallySignedTransaction& psbt, CTPSBTData& data);

// Assemble the kernel output script from fully-populated kernel data
// (including the signature).
bool BuildCTKernelScript(const CTKernelData& kernel, CScript& script);

// Enumerate the wallet's unspent confidential outputs (receipts that are
// not yet spent on-chain).  Returns (outpoint, receipt) pairs.
bool ListUnspentCTOutputs(const CWallet& wallet, const std::vector<COutPoint>& outpoints,
                          std::vector<std::pair<COutPoint, CTReceipt>>& out);

// Result of building a mint transaction (transparent value into CT).
struct MintCreationResult {
    CMutableTransaction tx;
    std::vector<CTReceipt> new_receipts; // one per output, same order
    int kernel_index{-1};
};

// Build a mint transaction: spends transparent inputs (satoshis) and
// creates confidential outputs of the given attosat amounts.  The input
// values enter the kernel balance with zero blinding, so the kernel
// signature only covers the output blindings; the transparent inputs are
// signed normally afterwards.  The explicit fee is in attosats.
util::Result<MintCreationResult> CreateMintTransaction(
    const std::vector<std::pair<COutPoint, CAmount>>& transparent_inputs,
    const std::vector<CAmount>& output_amounts, CAmount fee, FastRandomContext& rng);

// Build a fully confidential transaction.  Inputs are CT outputs the wallet
// already holds receipts for; outputs are new confidential outputs of the
// given attosat amounts.  The fee is in attosats.  Input witnesses carry
// the stored range proofs; the kernel carries the 128-bit fee and the
// excess signature.  Callers persist the returned receipts (keyed by the
// transaction id) before broadcasting.
util::Result<CTCreationResult> CreateConfidentialTransaction(
    const std::vector<std::pair<COutPoint, CTReceipt>>& ct_inputs,
    const std::vector<CAmount>& output_amounts, CAmount fee, FastRandomContext& rng);

} // namespace wallet

#endif // LIMENKA_WALLET_CT_H