// Package psbt writes the BIP174 form of a spend the wallet is about to make.
//
// The wallet signs locally, so this is not how a signature gets made; it is how a spend is
// described to anything else that speaks PSBT — the swap, a coordinator, a person checking the
// bytes. What it must get right is the record set: the witness UTXO and the witness script are
// what let a signer know what it is signing, and the SHA-256 preimage record is what lets a
// P2WSH claim be finalised without the preimage being in the transaction.
package psbt

import (
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"

	"git.smesh.lol/tinywallet/pkg/tx"
)

// Magic is the file magic: "psbt" and 0xff.
var magic = []byte{0x70, 0x73, 0x62, 0x74, 0xff}

// The record keys this package writes.
const (
	// globalUnsignedTx is the global map's only required key.
	globalUnsignedTx byte = 0x00
	// inputWitnessUTXO carries the output being spent, for segwit inputs.
	inputWitnessUTXO byte = 0x01
	// inputWitnessScript is the script of a P2WSH input.
	inputWitnessScript byte = 0x05
	// inputSHA256Preimage carries the preimage for an OP_SHA256 check.
	inputSHA256Preimage byte = 0x0b
)

// Errors the writer refuses.
var (
	// ErrNoInputs is a spend with nothing to spend.
	ErrNoInputs = errors.New("psbt: a psbt needs an input")
	// ErrLockedByPreimage is a preimage whose hash is not the one the script demands.
	ErrLockedByPreimage = errors.New("psbt: the preimage does not match the secret hash")
)

// SegwitInput describes one segwit input the wallet is spending.
type SegwitInput struct {
	// Amount is the value of the output being spent, in satoshis. BIP174 requires it for a
	// witness input, and BIP143 commits to it.
	Amount uint64
	// ScriptPubKey is the script being spent.
	ScriptPubKey []byte
	// WitnessScript is the P2WSH script, if this is a script-hash spend.
	WitnessScript []byte
	// Preimage is the SHA-256 preimage, if the script needs one.
	Preimage []byte
	// SecretHash is the hash the script compares the preimage against.
	SecretHash [32]byte
}

// Write renders a PSBT for an unsigned transaction and its inputs.
func Write(transaction *tx.Tx, inputs []SegwitInput) ([]byte, error) {
	if len(inputs) == 0 || len(inputs) != len(transaction.Inputs) {
		return nil, ErrNoInputs
	}

	unsigned, err := transaction.Serialize()

	if err != nil {
		return nil, err
	}

	out := append([]byte{}, magic...)
	out = appendRecord(out, globalUnsignedTx, unsigned)
	out = append(out, 0x00) // end of the global map

	for _, input := range inputs {
		utxo := make([]byte, 0, 8+len(input.ScriptPubKey))
		utxo = binary.LittleEndian.AppendUint64(utxo, input.Amount)
		utxo = append(utxo, compactSize(len(input.ScriptPubKey))...)
		utxo = append(utxo, input.ScriptPubKey...)

		out = appendRecord(out, inputWitnessUTXO, utxo)

		if len(input.WitnessScript) > 0 {
			out = appendRecord(out, inputWitnessScript, input.WitnessScript)
		}

		if len(input.Preimage) > 0 {
			if err := checkPreimage(input); err != nil {
				return nil, err
			}

			key := append([]byte{inputSHA256Preimage}, input.SecretHash[:]...)

			out = appendKeyedRecord(out, key, input.Preimage)
		}

		out = append(out, 0x00) // end of this input's map
	}

	for range transaction.Outputs {
		out = append(out, 0x00) // an output map, empty
	}

	return out, nil
}

// checkPreimage refuses a preimage that does not open the hash it is offered with.
func checkPreimage(input SegwitInput) error {
	sum := sha256Of(input.Preimage)

	if string(sum) != string(input.SecretHash[:]) {
		return ErrLockedByPreimage
	}

	return nil
}

// appendRecord writes a single-key record.
func appendRecord(out []byte, key byte, value []byte) []byte {
	return appendKeyedRecord(out, []byte{key}, value)
}

// appendKeyedRecord writes one key and its value in BIP174's form.
func appendKeyedRecord(out, key, value []byte) []byte {
	out = append(out, compactSize(len(key))...)
	out = append(out, key...)
	out = append(out, compactSize(len(value))...)

	return append(out, value...)
}

// compactSize writes BIP174's length prefix.
func compactSize(value int) []byte {
	switch {
	case value < 0xfd:
		return []byte{byte(value)}
	case value <= 0xffff:
		return []byte{0xfd, byte(value), byte(value >> 8)}
	}

	return []byte{0xfe, byte(value), byte(value >> 8), byte(value >> 16), byte(value >> 24)}
}

// sha256Of is the hash the HTLC's OP_SHA256 computes.
func sha256Of(preimage []byte) []byte {
	sum := sha256.Sum256(preimage)

	return sum[:]
}

// String names a record for an error message.
func recordName(key byte) string {
	switch key {
	case globalUnsignedTx:
		return "unsigned transaction"
	case inputWitnessUTXO:
		return "witness utxo"
	case inputWitnessScript:
		return "witness script"
	case inputSHA256Preimage:
		return "sha256 preimage"
	}

	return fmt.Sprintf("record %#x", key)
}