"""Bech32 encoding (BIP-173 variant) for nostr nsec/npub strings.

Reference implementation adapted from BIP-173 (Pieter Wuille, MIT licensed).
KEZ uses the original Bech32 checksum constant (not Bech32m), matching the
nostr NIP-19 convention and the Rust/Node implementations.
"""

from __future__ import annotations

CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"


def _polymod(values: list[int]) -> int:
    generator = [0x3B6A57B2, 0x26508E6D, 0x1EA119FA, 0x3D4233DD, 0x2A1462B3]
    chk = 1
    for v in values:
        b = chk >> 25
        chk = ((chk & 0x1FFFFFF) << 5) ^ v
        for i in range(5):
            chk ^= generator[i] if ((b >> i) & 1) else 0
    return chk


def _hrp_expand(hrp: str) -> list[int]:
    return [ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp]


def _verify_checksum(hrp: str, data: list[int]) -> bool:
    return _polymod(_hrp_expand(hrp) + data) == 1


def _create_checksum(hrp: str, data: list[int]) -> list[int]:
    values = _hrp_expand(hrp) + data
    polymod = _polymod(values + [0, 0, 0, 0, 0, 0]) ^ 1
    return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]


def _bech32_encode(hrp: str, data: list[int]) -> str:
    combined = data + _create_checksum(hrp, data)
    return hrp + "1" + "".join(CHARSET[d] for d in combined)


def _bech32_decode(bech: str) -> tuple[str, list[int]]:
    if any(ord(x) < 33 or ord(x) > 126 for x in bech):
        raise ValueError("bech32: invalid character")
    if bech.lower() != bech and bech.upper() != bech:
        raise ValueError("bech32: mixed case")
    bech = bech.lower()
    pos = bech.rfind("1")
    if pos < 1 or pos + 7 > len(bech):
        raise ValueError("bech32: invalid separator position")
    hrp = bech[:pos]
    if any(c not in CHARSET for c in bech[pos + 1 :]):
        raise ValueError("bech32: invalid data character")
    data = [CHARSET.find(c) for c in bech[pos + 1 :]]
    if not _verify_checksum(hrp, data):
        raise ValueError("bech32: bad checksum")
    return hrp, data[:-6]


def _convertbits(data: bytes | list[int], frombits: int, tobits: int, pad: bool) -> list[int]:
    acc = 0
    bits = 0
    ret: list[int] = []
    maxv = (1 << tobits) - 1
    max_acc = (1 << (frombits + tobits - 1)) - 1
    for value in data:
        if value < 0 or (value >> frombits):
            raise ValueError("bech32: invalid value in convertbits")
        acc = ((acc << frombits) | value) & max_acc
        bits += frombits
        while bits >= tobits:
            bits -= tobits
            ret.append((acc >> bits) & maxv)
    if pad:
        if bits:
            ret.append((acc << (tobits - bits)) & maxv)
    elif bits >= frombits or ((acc << (tobits - bits)) & maxv):
        raise ValueError("bech32: invalid padding in convertbits")
    return ret


def encode(hrp: str, payload: bytes) -> str:
    """Encode raw ``payload`` bytes as a bech32 string under ``hrp``."""
    data = _convertbits(payload, 8, 5, True)
    return _bech32_encode(hrp, data)


def decode(expected_hrp: str, bech: str) -> bytes:
    """Decode a bech32 string, asserting its HRP and returning the raw bytes."""
    hrp, data = _bech32_decode(bech)
    if hrp != expected_hrp:
        raise ValueError(f"bech32: expected hrp {expected_hrp!r}, got {hrp!r}")
    return bytes(_convertbits(data, 5, 8, False))