package psbt

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"testing"

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

// sha256Sum is the hash, named locally so the test does not depend on the implementation.
func sha256Sum(data []byte) [32]byte { return sha256.Sum256(data) }

// build makes an HTLC spend and the PSBT that describes it.
func build(t *testing.T) (*tx.Tx, SegwitInput, []byte) {
	t.Helper()

	preimage := make([]byte, 32)
	preimage[31] = 0x01

	var secretHash [32]byte
	copy(secretHash[:], sha256Of(preimage))

	claimant := bytes.Repeat([]byte{0x02}, 33)
	refunder := bytes.Repeat([]byte{0x03}, 33)

	witnessScript, err := script.HTLC(claimant, refunder, secretHash, 500)

	if err != nil {
		t.Fatal(err)
	}

	transaction := tx.New()

	var outpoint [32]byte
	outpoint[0] = 0x44

	transaction.AddInput(tx.Outpoint{Hash: outpoint, Index: 1}, tx.SequenceFinal)
	transaction.AddOutput(90000, script.P2WSH(witnessScript))

	input := SegwitInput{
		Amount:        100000,
		ScriptPubKey:  script.P2WSH(witnessScript),
		WitnessScript: witnessScript,
		Preimage:      preimage,
		SecretHash:    secretHash,
	}

	return transaction, input, witnessScript
}

// TestThePSBTIsWellFormed: magic, an unsigned transaction, one input map with the records a
// script-hash spend needs, and an empty output map.
func TestThePSBTIsWellFormed(t *testing.T) {
	transaction, input, witnessScript := build(t)

	raw, err := Write(transaction, []SegwitInput{input})

	if err != nil {
		t.Fatal(err)
	}

	if !bytes.HasPrefix(raw, magic) {
		t.Fatalf("the file starts %x", raw[:5])
	}

	unsigned, err := transaction.Serialize()

	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Contains(raw, unsigned) {
		t.Fatal("the unsigned transaction is not in the psbt")
	}

	if !bytes.Contains(raw, witnessScript) {
		t.Fatal("the witness script is not in the psbt")
	}

	if !bytes.Contains(raw, input.Preimage) {
		t.Fatal("the preimage is not in the psbt")
	}

	// The preimage record's key is the type byte followed by the hash.
	if !bytes.Contains(raw, append([]byte{inputSHA256Preimage}, input.SecretHash[:]...)) {
		t.Fatal("the preimage record is not keyed by its hash")
	}
}

// TestTheWitnessUTXOCarriesTheAmount: BIP174 and BIP143 both need the value.
func TestTheWitnessUTXOCarriesTheAmount(t *testing.T) {
	transaction, input, _ := build(t)

	raw, err := Write(transaction, []SegwitInput{input})

	if err != nil {
		t.Fatal(err)
	}

	found := false

	for start := 0; start+9 < len(raw); start++ {
		if raw[start] != 0x01 || binary.LittleEndian.Uint64(raw[start+2:start+10]) != input.Amount {
			continue
		}

		found = true

		break
	}

	if !found {
		t.Fatal("no witness utxo record carries the amount")
	}
}

// TestAWrongPreimageIsRefused: a preimage that does not open the hash is a mistake, and a
// PSBT carrying it would fail later, on chain, where it costs a fee.
func TestAWrongPreimageIsRefused(t *testing.T) {
	transaction, input, _ := build(t)

	input.Preimage = make([]byte, 32)

	if _, err := Write(transaction, []SegwitInput{input}); err != ErrLockedByPreimage {
		t.Fatalf("a wrong preimage was written: %v", err)
	}
}

// TestInputCountMustMatch: one description per input, or none.
func TestInputCountMustMatch(t *testing.T) {
	transaction, _, _ := build(t)

	if _, err := Write(transaction, nil); err != ErrNoInputs {
		t.Fatalf("no inputs: %v", err)
	}
}

// TestRecordNames: the names are what an error would say.
func TestRecordNames(t *testing.T) {
	if recordName(inputSHA256Preimage) != "sha256 preimage" {
		t.Fatal("the preimage record has no name")
	}
}

// TestWritingAndReadingAgree: what the writer produces, the reader understands, record for
// record — which is the only check that keeps the two directions honest.
func TestWritingAndReadingAgree(t *testing.T) {
	transaction, input, witnessScript := build(t)

	raw, err := Write(transaction, []SegwitInput{input})

	if err != nil {
		t.Fatal(err)
	}

	file, err := Read(raw)

	if err != nil {
		t.Fatal(err)
	}

	unsigned, err := transaction.Serialize()

	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(file.Unsigned, unsigned) {
		t.Fatal("the unsigned transaction did not survive")
	}

	if len(file.Inputs) != 1 {
		t.Fatalf("%d input maps", len(file.Inputs))
	}

	read := file.Inputs[0]

	if read.Amount != input.Amount || !bytes.Equal(read.ScriptPubKey, input.ScriptPubKey) {
		t.Fatalf("the witness utxo reads as %d %x", read.Amount, read.ScriptPubKey)
	}

	if !bytes.Equal(read.WitnessScript, witnessScript) {
		t.Fatal("the witness script did not survive")
	}

	if len(read.Preimages) != 1 || !bytes.Equal(read.Preimages[0].Preimage, input.Preimage) {
		t.Fatal("the preimage did not survive")
	}

	if read.Preimages[0].SecretHash != input.SecretHash {
		t.Fatal("the preimage is keyed by a different hash")
	}
}

// TestNotAPSBT: magic first, everything else after.
func TestNotAPSBT(t *testing.T) {
	if _, err := Read([]byte("not a psbt")); err != ErrNotPSBT {
		t.Fatalf("a non-psbt: %v", err)
	}
}

// TestATruncatedFileIsRefused: a length that runs past the end is an error, not a shorter file.
func TestATruncatedFileIsRefused(t *testing.T) {
	transaction, input, _ := build(t)

	raw, err := Write(transaction, []SegwitInput{input})

	if err != nil {
		t.Fatal(err)
	}

	for _, cut := range []int{6, 12, 40, len(raw) - 1} {
		if _, err := Read(raw[:cut]); err == nil {
			t.Fatalf("a %d-byte psbt parsed", cut)
		}
	}
}