// Package bech32 encodes and decodes segwit addresses.
//
// It implements only the part a swap wallet needs: P2WPKH and P2WSH, which are witness
// version 0 and therefore the original BIP173 checksum. A future witness version needs the
// BIP350 checksum instead, so this package refuses any version but 0 rather than encoding
// an address another wallet would reject.
package bech32

import (
	"fmt"
	"strings"
)

// charset is the bech32 alphabet, in the order the checksum depends on.
const charset = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"

// reverse maps a character to its value, built once.
var reverse = func() map[byte]int {
	out := make(map[byte]int, len(charset))

	for i := 0; i < len(charset); i++ {
		out[charset[i]] = i
	}

	return out
}()

// generator holds the bech32 checksum constants.
var generator = [5]uint32{0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3}

// polymod is the bech32 checksum function.
func polymod(values []int) uint32 {
	checksum := uint32(1)

	for _, value := range values {
		top := checksum >> 25
		checksum = (checksum&0x1ffffff)<<5 ^ uint32(value)

		for i := 0; i < 5; i++ {
			if (top>>uint(i))&1 == 1 {
				checksum ^= generator[i]
			}
		}
	}

	return checksum
}

// expandHRP spreads the human-readable part into the values the checksum covers.
func expandHRP(hrp string) []int {
	out := make([]int, 0, len(hrp)*2+1)

	for i := 0; i < len(hrp); i++ {
		out = append(out, int(hrp[i])>>5)
	}

	out = append(out, 0)

	for i := 0; i < len(hrp); i++ {
		out = append(out, int(hrp[i])&31)
	}

	return out
}

// convertBits regroups a byte slice into 5-bit groups, or back.
func convertBits(data []byte, from, to uint, pad bool) ([]int, error) {
	acc := uint32(0)
	bits := uint(0)
	maxValue := uint32(1)<<to - 1

	var out []int

	for _, value := range data {
		if uint32(value)>>from != 0 {
			return nil, fmt.Errorf("value %d does not fit in %d bits", value, from)
		}

		acc = acc<<from | uint32(value)
		bits += from

		for bits >= to {
			bits -= to
			out = append(out, int((acc>>bits)&maxValue))
		}
	}

	if pad {
		if bits > 0 {
			out = append(out, int((acc<<(to-bits))&maxValue))
		}

		return out, nil
	}

	if bits >= from {
		return nil, fmt.Errorf("excess padding")
	}

	if (acc<<(to-bits))&maxValue != 0 {
		return nil, fmt.Errorf("non-zero padding")
	}

	return out, nil
}

// Encode turns a human-readable part and a payload into a bech32 string with the BIP173
// checksum.
func Encode(hrp string, data []int) (string, error) {
	if hrp == "" {
		return "", fmt.Errorf("a bech32 string needs a human-readable part")
	}

	for i := 0; i < len(hrp); i++ {
		if hrp[i] < 33 || hrp[i] > 126 {
			return "", fmt.Errorf("the human-readable part contains %q", hrp[i])
		}

		if hrp[i] >= 'A' && hrp[i] <= 'Z' {
			return "", fmt.Errorf("the human-readable part must be lowercase")
		}
	}

	values := append(expandHRP(hrp), data...)
	values = append(values, 0, 0, 0, 0, 0, 0)

	checksum := polymod(values) ^ 1

	var out strings.Builder

	out.WriteString(hrp)
	out.WriteByte('1')

	for _, value := range data {
		out.WriteByte(charset[value])
	}

	for i := 0; i < 6; i++ {
		out.WriteByte(charset[(checksum>>uint(5*(5-i)))&31])
	}

	return out.String(), nil
}

// Decode splits a bech32 string into its human-readable part and payload values.
func Decode(text string) (string, []int, error) {
	if len(text) < 8 {
		return "", nil, fmt.Errorf("a bech32 string is at least 8 characters")
	}

	if len(text) > 90 {
		return "", nil, fmt.Errorf("a bech32 address is at most 90 characters, not %d",
			len(text))
	}

	// Mixed case is not valid: the checksum is defined over one case, and accepting both
	// would let two different strings be the same address.
	lower := strings.ToLower(text)

	if lower != text && strings.ToUpper(text) != text {
		return "", nil, fmt.Errorf("a bech32 string may not mix upper and lower case")
	}

	text = lower

	separator := strings.LastIndexByte(text, '1')

	if separator < 1 || separator+7 > len(text) {
		return "", nil, fmt.Errorf("a bech32 string needs a separator and a checksum")
	}

	hrp := text[:separator]
	data := make([]int, 0, len(text)-separator-1)

	for i := separator + 1; i < len(text); i++ {
		value, ok := reverse[text[i]]

		if !ok {
			return "", nil, fmt.Errorf("%q is not a bech32 character", text[i])
		}

		data = append(data, value)
	}

	if polymod(append(expandHRP(hrp), data...)) != 1 {
		return "", nil, fmt.Errorf("the bech32 checksum does not match")
	}

	return hrp, data[:len(data)-6], nil
}

// EncodeSegwit builds a witness version 0 address: P2WPKH for a 20-byte program, P2WSH for
// a 32-byte one.
func EncodeSegwit(hrp string, version byte, program []byte) (string, error) {
	if version != 0 {
		return "", fmt.Errorf(
			"this wallet encodes witness version 0 only; version %d needs the BIP350 "+
				"checksum", version,
		)
	}

	if len(program) != 20 && len(program) != 32 {
		return "", fmt.Errorf(
			"a witness version 0 program is 20 or 32 bytes, not %d", len(program),
		)
	}

	grouped, err := convertBits(program, 8, 5, true)

	if err != nil {
		return "", err
	}

	return Encode(hrp, append([]int{int(version)}, grouped...))
}

// DecodeSegwit reads a witness address and returns its version and program.
func DecodeSegwit(address string) (string, byte, []byte, error) {
	hrp, data, err := Decode(address)

	if err != nil {
		return "", 0, nil, err
	}

	if len(data) < 1 {
		return "", 0, nil, fmt.Errorf("the address has no witness version")
	}

	if data[0] != 0 {
		return "", 0, nil, fmt.Errorf(
			"this wallet reads witness version 0 only; the address is version %d", data[0],
		)
	}

	program, err := convertBits(intsToBytes(data[1:]), 5, 8, false)

	if err != nil {
		return "", 0, nil, err
	}

	if len(program) != 20 && len(program) != 32 {
		return "", 0, nil, fmt.Errorf("the witness program is %d bytes", len(program))
	}

	out := make([]byte, len(program))

	for i, value := range program {
		out[i] = byte(value)
	}

	return hrp, 0, out, nil
}

// intsToBytes narrows the 5-bit values into a byte slice.
func intsToBytes(values []int) []byte {
	out := make([]byte, len(values))

	for i, value := range values {
		out[i] = byte(value)
	}

	return out
}