package helpers

import "git.smesh.lol/moxie/pkg/mxutil"

// Bech32 encoding/decoding for NIP-19.
// Implements bech32 (BIP-173) without external deps.

const bech32Charset = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"

// Bech32Encode encodes data with the given human-readable part. The output
// size is known once the checksum exists, so it is presized and written
// through a cursor: `buf | "1"` had a non-empty left side, and concat then
// allocated a fresh exact-size buffer, leaving the first value push no room.
func Bech32Encode(hrp string, data []byte) (s string) {
	values := bytesToBase32(data)
	checksum := bech32Checksum(hrp, values)
	values = values | checksum

	buf := []byte{:len(hrp) + 1 + len(values)}
	j := 0
	for i := 0; i < len(hrp); i++ {
		buf[j] = hrp[i]
		j++
	}
	buf[j] = '1'
	j++
	for _, v := range values {
		buf[j] = bech32Charset[v]
		j++
	}
	return string(buf)
}

// Bech32Decode decodes a bech32 string. Returns hrp and data bytes.
func Bech32Decode(s string) (hrp string, data []byte) {
	// Find separator.
	pos := -1
	for i := len(s) - 1; i >= 0; i-- {
		if s[i] == '1' {
			pos = i
			break
		}
	}
	if pos < 1 || pos+7 > len(s) {
		return "", nil
	}

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

	values := []byte{:len(dataStr)}
	for i := 0; i < len(dataStr); i++ {
		idx := charsetIndex(dataStr[i])
		if idx < 0 {
			return "", nil
		}
		values[i] = byte(idx)
	}

	if !bech32Verify(hrp, values) {
		return "", nil
	}

	// Strip checksum (last 6 chars).
	values = values[:len(values)-6]
	data := base32ToBytes(values)
	return hrp, data
}

// NIP-19 helpers.

// EncodeNpub encodes a 32-byte public key as npub.
func EncodeNpub(pubkey []byte) (s string) {
	return Bech32Encode("npub", pubkey)
}

// EncodeNsec encodes a 32-byte secret key as nsec.
func EncodeNsec(seckey []byte) (s string) {
	return Bech32Encode("nsec", seckey)
}

// EncodeNote encodes a 32-byte event ID as note.
func EncodeNote(eventID []byte) (s string) {
	return Bech32Encode("note", eventID)
}

// EncodeNevent encodes an event reference as nevent (NIP-19 TLV).
// The TLV length is a function of its parts, so it is presized exactly: push
// does not grow, and the `data | idBytes` step returns a full slice.
func EncodeNevent(id string, relays []string, author string) (s string) {
	idBytes := HexDecode(id)
	ab := HexDecode(author)
	hasID := len(idBytes) == 32
	hasAuthor := len(ab) == 32
	n := 0
	if hasID {
		n += 34
	}
	for _, r := range relays {
		n += 2 + len(r)
	}
	if hasAuthor {
		n += 34
	}
	data := []byte{:n}
	j := 0
	if hasID {
		data[j] = 0
		data[j+1] = 32
		j += 2
		for _, b := range idBytes {
			data[j] = b
			j++
		}
	}
	for _, r := range relays {
		data[j] = 1
		data[j+1] = byte(len(r))
		j += 2
		for k := 0; k < len(r); k++ {
			data[j] = r[k]
			j++
		}
	}
	if hasAuthor {
		data[j] = 2
		data[j+1] = 32
		j += 2
		for _, b := range ab {
			data[j] = b
			j++
		}
	}
	return Bech32Encode("nevent", data[:j])
}

// DecodeNpub decodes an npub string to 32 bytes.
func DecodeNpub(s string) (buf []byte) {
	hrp, data := Bech32Decode(s)
	if hrp != "npub" || len(data) != 32 {
		return nil
	}
	return data
}

// DecodeNsec decodes an nsec string to 32 bytes.
func DecodeNsec(s string) (buf []byte) {
	hrp, data := Bech32Decode(s)
	if hrp != "nsec" || len(data) != 32 {
		return nil
	}
	return data
}

// DecodeNote decodes a note string to 32 bytes.
func DecodeNote(s string) (buf []byte) {
	hrp, data := Bech32Decode(s)
	if hrp != "note" || len(data) != 32 {
		return nil
	}
	return data
}

// Nevent holds decoded nevent TLV data (NIP-19).
type Nevent struct {
	ID     string   // hex event ID
	Relays []string // optional relay hints
	Author string   // hex pubkey (optional)
}

// DecodeNevent decodes a nevent1... bech32 string (TLV format).
func DecodeNevent(s string) (n *Nevent) {
	hrp, data := Bech32Decode(s)
	if hrp != "nevent" {
		return nil
	}
	result := &Nevent{}
	i := 0
	for i+2 <= len(data) {
		t := data[i]
		l := int32(data[i+1])
		i += 2
		if i+l > len(data) {
			break
		}
		v := data[i : i+l]
		i += l
		switch t {
		case 0:
			if l == 32 {
				result.ID = HexEncode(v)
			}
		case 1:
			result.Relays = mxutil.Ensure(result.Relays, 1)
			result.Relays = push(result.Relays, string(v))
		case 2:
			if l == 32 {
				result.Author = HexEncode(v)
			}
		}
	}
	if result.ID == "" {
		return nil
	}
	return result
}

// Nprofile holds decoded nprofile TLV data (NIP-19).
type Nprofile struct {
	Pubkey string   // hex pubkey
	Relays []string // optional relay hints
}

// DecodeNprofile decodes an nprofile1... bech32 string (TLV format).
func DecodeNprofile(s string) (n *Nprofile) {
	hrp, data := Bech32Decode(s)
	if hrp != "nprofile" {
		return nil
	}
	result := &Nprofile{}
	i := 0
	for i+2 <= len(data) {
		t := data[i]
		l := int32(data[i+1])
		i += 2
		if i+l > len(data) {
			break
		}
		v := data[i : i+l]
		i += l
		switch t {
		case 0:
			if l == 32 {
				result.Pubkey = HexEncode(v)
			}
		case 1:
			result.Relays = mxutil.Ensure(result.Relays, 1)
			result.Relays = push(result.Relays, string(v))
		}
	}
	if result.Pubkey == "" {
		return nil
	}
	return result
}

// Naddr holds decoded naddr TLV data (NIP-19).
type Naddr struct {
	Kind   uint32
	Pubkey string   // hex pubkey
	D      string   // d-tag identifier
	Relays []string // optional relay hints
}

// DecodeNaddr decodes an naddr1... bech32 string (TLV format).
func DecodeNaddr(s string) (n *Naddr) {
	hrp, data := Bech32Decode(s)
	if hrp != "naddr" {
		return nil
	}
	result := &Naddr{}
	i := 0
	for i+2 <= len(data) {
		t := data[i]
		l := int32(data[i+1])
		i += 2
		if i+l > len(data) {
			break
		}
		v := data[i : i+l]
		i += l
		switch t {
		case 0:
			result.D = string(v)
		case 1:
			result.Relays = mxutil.Ensure(result.Relays, 1)
			result.Relays = push(result.Relays, string(v))
		case 2:
			if l == 32 {
				result.Pubkey = HexEncode(v)
			}
		case 3:
			if l == 4 {
				result.Kind = uint32(v[0])<<24 | uint32(v[1])<<16 | uint32(v[2])<<8 | uint32(v[3])
			}
		}
	}
	if result.Pubkey == "" || result.Kind == 0 {
		return nil
	}
	return result
}

// NaddrCoordKey returns the canonical "kind:pubkey:d" lookup key.
func NaddrCoordKey(kind uint32, pubkey, d string) (s string) {
	k := ""
	n := kind
	if n == 0 {
		k = "0"
	} else {
		for n > 0 {
			k = string([]byte{byte('0' + n%10)}) | k
			n /= 10
		}
	}
	return k | ":" | pubkey | ":" | d
}

// PubkeyShort returns first 8 chars of hex pubkey.
func PubkeyShort(pubkey string) (s string) {
	if len(pubkey) >= 8 {
		return pubkey[:8]
	}
	return pubkey
}

// Internal bech32 functions.

// bytesToBase32 converts 8-bit bytes to 5-bit groups. The output length is a
// function of the input length, so it is presized and written through a
// cursor: push does not grow, and a nil sink stops at four values.
func bytesToBase32(data []byte) (buf []byte) {
	out := []byte{:(len(data)*8 + 4) / 5}
	j := 0
	acc := 0
	bits := 0
	for _, b := range data {
		acc = (acc << 8) | int32(b)
		bits += 8
		for bits >= 5 {
			bits -= 5
			out[j] = byte((acc >> bits) & 0x1f)
			j++
		}
		acc &= (1 << uint32(bits)) - 1
	}
	if bits > 0 {
		out[j] = byte((acc << (5 - bits)) & 0x1f)
		j++
	}
	return out[:j]
}

func base32ToBytes(data []byte) (buf []byte) {
	out := []byte{:len(data) * 5 / 8}
	j := 0
	acc := 0
	bits := 0
	for _, v := range data {
		acc = (acc << 5) | int32(v)
		bits += 5
		for bits >= 8 {
			bits -= 8
			out[j] = byte((acc >> bits) & 0xff)
			j++
		}
		acc &= (1 << uint32(bits)) - 1
	}
	return out[:j]
}

func bech32Polymod(values []byte) (n uint32) {
	gen := [5]uint32{0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3}
	chk := uint32(1)
	for _, v := range values {
		b := chk >> 25
		chk = ((chk & 0x1ffffff) << 5) ^ uint32(v)
		for i := 0; i < 5; i++ {
			if (b>>uint32(i))&1 == 1 {
				chk ^= gen[i]
			}
		}
	}
	return chk
}

func bech32HRPExpand(hrp string) (buf []byte) {
	out := []byte{:0:len(hrp)*2+1}
	for i := 0; i < len(hrp); i++ {
		out = push(out, byte(hrp[i]>>5))
	}
	out = push(out, 0)
	for i := 0; i < len(hrp); i++ {
		out = push(out, byte(hrp[i]&0x1f))
	}
	return out
}

func bech32Checksum(hrp string, data []byte) (buf []byte) {
	exp := bech32HRPExpand(hrp)
	values := []byte{:len(exp) + len(data) + 6}
	j := 0
	for _, b := range exp {
		values[j] = b
		j++
	}
	for _, b := range data {
		values[j] = b
		j++
	}
	// Six zero bytes pad the checksum input (BIP-173).
	for i := 0; i < 6; i++ {
		values[j] = 0
		j++
	}
	polymod := bech32Polymod(values) ^ 1
	out := []byte{:6}
	for i := 0; i < 6; i++ {
		out[i] = byte((polymod >> (5 * (5 - uint32(i)))) & 0x1f)
	}
	return out
}

func bech32Verify(hrp string, data []byte) (ok bool) {
	values := bech32HRPExpand(hrp) | data
	return bech32Polymod(values) == 1
}

func charsetIndex(c byte) (n int32) {
	for i := 0; i < len(bech32Charset); i++ {
		if bech32Charset[i] == c {
			return i
		}
	}
	return -1
}