package secp256k1

import (
	"fmt"
	"math/big"
	"strings"
)

// Bech32 alphabet (no 1, b, i, o to avoid confusion)
const bech32Alphabet = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"

// bech32Polymod computes the Bech32 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
func bech32HRPExpand(hrp string) []int {
	result := make([]int, len(hrp)*2+1)
	for i, c := range hrp {
		result[i] = int(c) >> 5
	}
	result[len(hrp)] = 0
	for i, c := range hrp {
		result[len(hrp)+1+i] = int(c) & 31
	}
	return result
}

// bech32CreateChecksum creates a 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 a bech32 checksum
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 := []int{}
	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 a human-readable prefix
func Bech32Encode(hrp string, data []byte) (string, error) {
	// Convert 8-bit bytes to 5-bit groups
	values, err := convertBits(data, 8, 5, true)
	if err != nil {
		return "", err
	}

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

	// Build result string
	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
func Bech32Decode(s string) (string, []byte, error) {
	// Find separator
	pos := strings.LastIndex(s, "1")
	if pos < 1 || pos+7 > len(s) {
		return "", nil, fmt.Errorf("invalid bech32: separator position")
	}

	// Split HRP and data
	hrp := strings.ToLower(s[:pos])
	dataStr := strings.ToLower(s[pos+1:])

	// Decode data characters
	data := make([]int, len(dataStr))
	for i, c := range dataStr {
		idx := strings.IndexRune(bech32Alphabet, c)
		if idx < 0 {
			return "", nil, fmt.Errorf("invalid bech32: character %c", c)
		}
		data[i] = idx
	}

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

	// Remove checksum and convert back to 8-bit
	values := data[:len(data)-6]
	bytes, err := convertBitsToBytes(values)
	if err != nil {
		return "", nil, err
	}

	return hrp, bytes, nil
}

// convertBitsToBytes converts 5-bit groups back to bytes
func convertBitsToBytes(data []int) ([]byte, error) {
	acc := 0
	bits := 0
	result := []byte{}

	for _, value := range data {
		acc = (acc << 5) | value
		bits += 5
		for bits >= 8 {
			bits -= 8
			result = append(result, byte((acc>>bits)&0xff))
		}
	}

	return result, nil
}

// === Nostr-specific helpers ===

// Nsec returns the private key as a Nostr nsec string
func (priv *PrivateKey) Nsec() string {
	encoded, _ := Bech32Encode("nsec", priv.Bytes())
	return encoded
}

// Npub returns the public key as a Nostr npub string (x-only)
func (pub *PublicKey) Npub() string {
	encoded, _ := Bech32Encode("npub", pub.XOnlyBytes())
	return encoded
}

// PrivateKeyFromNsec parses an nsec string
func PrivateKeyFromNsec(nsec string) (*PrivateKey, error) {
	hrp, data, err := Bech32Decode(nsec)
	if err != nil {
		return nil, fmt.Errorf("invalid nsec: %w", err)
	}
	if hrp != "nsec" {
		return nil, fmt.Errorf("invalid nsec: wrong prefix %q", hrp)
	}
	if len(data) != 32 {
		return nil, fmt.Errorf("invalid nsec: wrong length %d", len(data))
	}
	return NewPrivateKeyFromBytes(data)
}

// PublicKeyFromNpub parses an npub string
func PublicKeyFromNpub(npub string) (*PublicKey, error) {
	hrp, data, err := Bech32Decode(npub)
	if err != nil {
		return nil, fmt.Errorf("invalid npub: %w", err)
	}
	if hrp != "npub" {
		return nil, fmt.Errorf("invalid npub: wrong prefix %q", hrp)
	}
	if len(data) != 32 {
		return nil, fmt.Errorf("invalid npub: wrong length %d", len(data))
	}

	// Lift x-coordinate to full point
	point, err := liftX(new(big.Int).SetBytes(data))
	if err != nil {
		return nil, fmt.Errorf("invalid npub: %w", err)
	}

	return &PublicKey{Point: point}, nil
}