"""Crockford base-32 encoding with double-SHA256 checksum (c32check).

Implements the c32check encoding used by Stacks addresses.
Reference: https://github.com/stacks-network/c32check
"""

from __future__ import annotations

import hashlib

C32_ALPHABET = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"

# Lookup table for decoding
_C32_DECODE = {c: i for i, c in enumerate(C32_ALPHABET)}
# Also accept lowercase
_C32_DECODE.update({c.lower(): i for i, c in enumerate(C32_ALPHABET)})


def c32_encode(data: bytes) -> str:
    """Encode raw bytes into a c32 string (no checksum)."""
    if not data:
        return ""

    # Convert bytes to a big integer
    num = int.from_bytes(data, "big")

    if num == 0:
        # Preserve leading zeros
        result = C32_ALPHABET[0] * len(data)
        return result

    # Base-32 encode
    chars: list[str] = []
    while num > 0:
        num, remainder = divmod(num, 32)
        chars.append(C32_ALPHABET[remainder])

    # Reverse to get most significant first
    chars.reverse()

    # Preserve leading zero bytes
    leading_zeros = 0
    for b in data:
        if b == 0:
            leading_zeros += 1
        else:
            break

    return C32_ALPHABET[0] * leading_zeros + "".join(chars)


def c32_decode(c32_str: str) -> bytes:
    """Decode a c32 string back to raw bytes (no checksum)."""
    if not c32_str:
        return b""

    # Normalize: uppercase, handle O→0, L/I→1
    normalized = c32_str.upper()
    normalized = normalized.replace("O", "0").replace("L", "1").replace("I", "1")

    # Count leading zeros
    leading_zeros = 0
    for c in normalized:
        if c == C32_ALPHABET[0]:
            leading_zeros += 1
        else:
            break

    # Convert from base-32 to integer
    num = 0
    for c in normalized:
        if c not in _C32_DECODE:
            raise ValueError(f"Invalid c32 character: {c}")
        num = num * 32 + _C32_DECODE[c]

    if num == 0:
        return b"\x00" * leading_zeros

    # Convert integer to bytes
    byte_length = (num.bit_length() + 7) // 8
    result = num.to_bytes(byte_length, "big")

    return b"\x00" * leading_zeros + result


def _double_sha256(data: bytes) -> bytes:
    """Compute double SHA-256 hash."""
    return hashlib.sha256(hashlib.sha256(data).digest()).digest()


def c32_check_encode(version: int, data: bytes) -> str:
    """Encode data with version byte and checksum using c32check.

    Format: version_byte + data + checksum[0:4]
    The checksum is double-SHA256 of (version_byte + data).
    """
    if version < 0 or version > 31:
        raise ValueError(f"Version must be 0-31, got {version}")

    # Compute checksum: double-SHA256(version + data)
    check_data = bytes([version]) + data
    checksum = _double_sha256(check_data)[:4]

    # Encode: data + checksum (version is separate)
    payload = data + checksum
    c32_str = c32_encode(payload)

    # Prefix with version character
    return C32_ALPHABET[version] + c32_str


def c32_check_decode(c32_check_str: str) -> tuple[int, bytes]:
    """Decode a c32check-encoded string, returning (version, data).

    Verifies the checksum and raises ValueError if invalid.
    """
    if len(c32_check_str) < 2:
        raise ValueError("c32check string too short")

    # First character is version
    version_char = c32_check_str[0].upper()
    version_char = version_char.replace("O", "0").replace("L", "1").replace("I", "1")

    if version_char not in _C32_DECODE:
        raise ValueError(f"Invalid version character: {c32_check_str[0]}")

    version = _C32_DECODE[version_char]

    # Decode the rest
    payload = c32_decode(c32_check_str[1:])

    if len(payload) < 4:
        raise ValueError("c32check data too short for checksum")

    # Split into data and checksum
    data = payload[:-4]
    checksum = payload[-4:]

    # Verify checksum
    expected_checksum = _double_sha256(bytes([version]) + data)[:4]
    if checksum != expected_checksum:
        raise ValueError("c32check checksum mismatch")

    return version, data


def c32_address(version: int, hash160_bytes: bytes) -> str:
    """Create a Stacks c32check address from version and hash160.

    Returns address prefixed with 'S'.
    """
    if len(hash160_bytes) != 20:
        raise ValueError(f"hash160 must be 20 bytes, got {len(hash160_bytes)}")

    c32_str = c32_check_encode(version, hash160_bytes)
    return "S" + c32_str


def c32_address_decode(address: str) -> tuple[int, bytes]:
    """Decode a Stacks c32check address, returning (version, hash160).

    Address must start with 'S'.
    """
    if not address or address[0] not in ("S", "s"):
        raise ValueError("Stacks address must start with 'S'")

    return c32_check_decode(address[1:])
