package script

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

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

// TestP2WPKHScript: OP_0, a twenty-byte push, the program.
func TestP2WPKHScript(t *testing.T) {
	var program [20]byte
	program[0], program[19] = 0xaa, 0xbb

	out := P2WPKH(program)

	if len(out) != 22 || out[0] != OpZero || out[1] != 0x14 || out[2] != 0xaa || out[21] != 0xbb {
		t.Fatalf("the script is %x", out)
	}
}

// TestP2WSHScript: OP_0, a thirty-two-byte push, the script hash.
func TestP2WSHScript(t *testing.T) {
	witness := []byte{OpIf, 0x01, 0x01, OpEndIf}

	out := P2WSH(witness)

	if len(out) != 34 || out[0] != OpZero || out[1] != 0x20 {
		t.Fatalf("the script is %x", out)
	}

	sum := sha256.Sum256(witness)

	if !bytes.Equal(out[2:], sum[:]) {
		t.Fatalf("the program is %x, want %x", out[2:], sum)
	}
}

// TestPushIsMinimal: the push width is what makes the script hash what it is.
func TestPushIsMinimal(t *testing.T) {
	cases := []struct {
		length int
		want   []byte
	}{
		{0, []byte{OpZero}},
		{1, []byte{0x01}},
		{75, []byte{75}},
		{76, []byte{0x4c, 76}},
		{0xff, []byte{0x4c, 0xff}},
		{0x100, []byte{0x4d, 0x00, 0x01}},
	}

	for _, test := range cases {
		out, err := Push(make([]byte, test.length))

		if err != nil {
			t.Fatalf("%d bytes: %v", test.length, err)
		}

		if !bytes.HasPrefix(out, test.want) {
			t.Fatalf("%d bytes: the prefix is %x, want %x", test.length, out[:len(test.want)], test.want)
		}

		if len(out) != len(test.want)+test.length {
			t.Fatalf("%d bytes: the push is %d bytes", test.length, len(out))
		}
	}
}

// TestAddressRoundTrip: an address is a program with a prefix, and decoding gives it back.
func TestAddressRoundTrip(t *testing.T) {
	var program [20]byte
	program[3] = 0x42

	address, err := P2WPKHAddress("bcrt", program)

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

	hrp, version, decoded, err := bech32.DecodeSegwit(address)

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

	if hrp != "bcrt" || version != 0 || !bytes.Equal(decoded, program[:]) {
		t.Fatalf("decoded %q %d %x", hrp, version, decoded)
	}

	witness := []byte{OpIf, OpEndIf}
	scriptAddress, err := P2WSHAddress("bcrt", witness)

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

	hrp, version, decoded, err = bech32.DecodeSegwit(scriptAddress)

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

	if hrp != "bcrt" || version != 0 || len(decoded) != 32 {
		t.Fatalf("decoded %q %d %x", hrp, version, decoded)
	}
}

// TestTheWitnessScriptDeterminesTheAddress: a one-byte difference is a different wallet.
func TestTheWitnessScriptDeterminesTheAddress(t *testing.T) {
	first, err := P2WSHAddress("bcrt", []byte{OpIf, OpOne, OpEndIf})

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

	second, err := P2WSHAddress("bcrt", []byte{OpIf, OpZero, OpEndIf})

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

	if first == second {
		t.Fatal("two different witness scripts produced one address")
	}
}

// TestP2WPKHMatchesTheWallet: the wallet's own address script and this one agree.
func TestP2WPKHMatchesTheWallet(t *testing.T) {
	var program [20]byte
	program[0] = 0x01

	address, err := P2WPKHAddress("bc", program)

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

	hrp, _, decoded, err := bech32.DecodeSegwit(address)

	if err != nil || hrp != "bc" || !bytes.Equal(decoded, program[:]) {
		t.Fatalf("decoded %q %x, %v", hrp, decoded, err)
	}
}