package nostr

import (
	"fmt"
	"strings"
)

// Bech32 encoding/decoding for NIP-19 (npub, nsec, note, etc.)
// Implements BIP-173 bech32 encoding.

const bech32Alphabet = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"

var bech32AlphabetMap [256]int8

func init() {
	for i := range bech32AlphabetMap {
		bech32AlphabetMap[i] = -1
	}
	for i, c := range bech32Alphabet {
		bech32AlphabetMap[c] = int8(i)
	}
}

// bech32Polymod computes the BCH checksum.
func bech32Polymod(values []int) int {
	gen := []int{0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3}
	chk := 1
	for _, v := range values {
		top := chk >> 25
		chk = (chk&0x1ffffff)<<5 ^ v
		for i := 0; i < 5; i++ {
			if (top>>i)&1 == 1 {
				chk ^= gen[i]
			}
		}
	}
	return chk
}

// bech32HRPExpand expands the human-readable part for checksum computation.
func bech32HRPExpand(hrp string) []int {
	result := make([]int, len(hrp)*2+1)
	for i, c := range hrp {
		result[i] = int(c) >> 5
		result[i+len(hrp)+1] = int(c) & 31
	}
	return result
}

// bech32CreateChecksum creates the 6-character checksum.
func bech32CreateChecksum(hrp string, data []int) []int {
	values := append(bech32HRPExpand(hrp), data...)
	values = append(values, []int{0, 0, 0, 0, 0, 0}...)
	polymod := bech32Polymod(values) ^ 1
	checksum := make([]int, 6)
	for i := 0; i < 6; i++ {
		checksum[i] = (polymod >> (5 * (5 - i))) & 31
	}
	return checksum
}

// bech32VerifyChecksum verifies the checksum of bech32 data.
func bech32VerifyChecksum(hrp string, data []int) bool {
	return bech32Polymod(append(bech32HRPExpand(hrp), data...)) == 1
}

// convertBits converts between bit groups.
func convertBits(data []byte, fromBits, toBits int, pad bool) ([]int, error) {
	acc := 0
	bits := 0
	result := make([]int, 0, len(data)*fromBits/toBits+1)
	maxv := (1 << toBits) - 1

	for _, value := range data {
		acc = (acc << fromBits) | int(value)
		bits += fromBits
		for bits >= toBits {
			bits -= toBits
			result = append(result, (acc>>bits)&maxv)
		}
	}

	if pad {
		if bits > 0 {
			result = append(result, (acc<<(toBits-bits))&maxv)
		}
	} else if bits >= fromBits || ((acc<<(toBits-bits))&maxv) != 0 {
		return nil, fmt.Errorf("invalid padding")
	}

	return result, nil
}

// Bech32Encode encodes data with the given human-readable prefix.
func Bech32Encode(hrp string, data []byte) (string, error) {
	values, err := convertBits(data, 8, 5, true)
	if err != nil {
		return "", err
	}

	checksum := bech32CreateChecksum(hrp, values)
	combined := append(values, checksum...)

	var result strings.Builder
	result.WriteString(hrp)
	result.WriteByte('1')
	for _, v := range combined {
		result.WriteByte(bech32Alphabet[v])
	}

	return result.String(), nil
}

// Bech32Decode decodes a bech32 string, returning the HRP and data.
func Bech32Decode(s string) (string, []byte, error) {
	s = strings.ToLower(s)

	pos := strings.LastIndexByte(s, '1')
	if pos < 1 || pos+7 > len(s) {
		return "", nil, fmt.Errorf("invalid bech32 string")
	}

	hrp := s[:pos]
	dataStr := s[pos+1:]

	data := make([]int, len(dataStr))
	for i, c := range dataStr {
		val := bech32AlphabetMap[c]
		if val == -1 {
			return "", nil, fmt.Errorf("invalid character: %c", c)
		}
		data[i] = int(val)
	}

	if !bech32VerifyChecksum(hrp, data) {
		return "", nil, fmt.Errorf("invalid checksum")
	}

	// Remove checksum
	data = data[:len(data)-6]

	// Convert from 5-bit to 8-bit
	result, err := convertBits(intSliceToBytes(data), 5, 8, false)
	if err != nil {
		return "", nil, err
	}

	bytes := make([]byte, len(result))
	for i, v := range result {
		bytes[i] = byte(v)
	}

	return hrp, bytes, nil
}

func intSliceToBytes(data []int) []byte {
	result := make([]byte, len(data))
	for i, v := range data {
		result[i] = byte(v)
	}
	return result
}