#!/usr/bin/env python3
# Copyright (c) 2017 Pieter Wuille, Shammah Chancellor, Neil Booth
#
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
#
# Originally taken from Electron-Cash:
# https://raw.githubusercontent.com/fyookball/electrum/master/lib/cashaddr.py
#
_CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"


def _polymod(values):
    """Internal function that computes the cashaddr checksum."""
    c = 1
    for d in values:
        c0 = c >> 35
        c = ((c & 0x07FFFFFFFF) << 5) ^ d
        if c0 & 0x01:
            c ^= 0x98F2BC8E61
        if c0 & 0x02:
            c ^= 0x79B76D99E2
        if c0 & 0x04:
            c ^= 0xF33E5FB3C4
        if c0 & 0x08:
            c ^= 0xAE2EABE2A8
        if c0 & 0x10:
            c ^= 0x1E4F43E470
    retval = c ^ 1
    return retval


def _prefix_expand(prefix):
    """Expand the prefix into values for checksum computation."""
    retval = bytearray(ord(x) & 0x1F for x in prefix)
    # Append null separator
    retval.append(0)
    return retval


def _create_checksum(prefix, data):
    """Compute the checksum values given prefix and data."""
    values = _prefix_expand(prefix) + data + bytes(8)
    polymod = _polymod(values)
    # Return the polymod expanded into eight 5-bit elements
    return bytes((polymod >> 5 * (7 - i)) & 31 for i in range(8))


def _convertbits(data, frombits, tobits, pad=True):
    """General power-of-2 base conversion."""
    acc = 0
    bits = 0
    ret = bytearray()
    maxv = (1 << tobits) - 1
    max_acc = (1 << (frombits + tobits - 1)) - 1
    for value in data:
        acc = ((acc << frombits) | value) & max_acc
        bits += frombits
        while bits >= tobits:
            bits -= tobits
            ret.append((acc >> bits) & maxv)

    if pad and bits:
        ret.append((acc << (tobits - bits)) & maxv)

    return ret


def _pack_addr_data(kind, addr_hash):
    """Pack addr data with version byte"""
    version_byte = kind << 3

    offset = 1
    encoded_size = 0
    if len(addr_hash) >= 40:
        offset = 2
        encoded_size |= 0x04
    encoded_size |= (len(addr_hash) - 20 * offset) // (4 * offset)

    # invalid size?
    if (len(addr_hash) - 20 * offset) % (4 * offset) != 0 or not 0 <= encoded_size <= 7:
        raise ValueError(f"invalid address hash size {addr_hash}")

    version_byte |= encoded_size

    data = bytes([version_byte]) + addr_hash
    return _convertbits(data, 8, 5, True)


def _decode_payload(addr):
    """Validate a cashaddr string.

    Throws CashAddr.Error if it is invalid, otherwise returns the
    tuple

       (prefix,  payload)

    without the checksum.
    """
    lower = addr.lower()
    if lower != addr and addr.upper() != addr:
        raise ValueError(f"mixed case in address: {addr}")

    parts = lower.split(":", 1)
    if len(parts) != 2:
        raise ValueError(f"address missing ':' separator: {addr}")

    prefix, payload = parts
    if not prefix:
        raise ValueError(f"address prefix is missing: {addr}")
    if not all(33 <= ord(x) <= 126 for x in prefix):
        raise ValueError(f"invalid address prefix: {prefix}")
    if len(payload) < 8 or len(payload) > 124:
        raise ValueError(f"address payload has invalid length: {len(addr)}")
    try:
        data = bytes(_CHARSET.find(x) for x in payload)
    except ValueError as e:
        raise ValueError(f"invalid characters in address: {payload}") from e

    if _polymod(_prefix_expand(prefix) + data):
        raise ValueError(f"invalid checksum in address: {addr}")

    if lower != addr:
        prefix = prefix.upper()

    # Drop the 40 bit checksum
    return prefix, data[:-8]


#
# External Interface
#


PUBKEY_TYPE = 0
SCRIPT_TYPE = 1
# cashtoken aware
PUBKEY_TYPE_TOKEN = 2
SCRIPT_TYPE_TOKEN = 3


def decode(address):
    """Given a cashaddr address, return a tuple

    (prefix, kind, hash)
    """
    if not isinstance(address, str):
        raise TypeError("address must be a string")

    prefix, payload = _decode_payload(address)

    # Ensure there isn't extra padding
    extrabits = len(payload) * 5 % 8
    if extrabits >= 5:
        raise ValueError(f"excess padding in address {address}")

    # Ensure extrabits are zeros
    if payload[-1] & ((1 << extrabits) - 1):
        raise ValueError(f"non-zero padding in address {address}")

    decoded = _convertbits(payload, 5, 8, False)
    version = decoded[0]
    addr_hash = bytes(decoded[1:])
    size = (version & 0x03) * 4 + 20
    # Double the size, if the 3rd bit is on.
    if version & 0x04:
        size <<= 1
    if size != len(addr_hash):
        raise ValueError(
            f"address hash has length {len(addr_hash)} but expected {size}"
        )

    kind = version >> 3
    if kind not in (SCRIPT_TYPE, PUBKEY_TYPE, SCRIPT_TYPE_TOKEN, PUBKEY_TYPE_TOKEN):
        raise ValueError(f"unrecognised address type {kind}")

    return prefix, kind, addr_hash


def encode(prefix, kind, addr_hash):
    """Encode a cashaddr address without prefix and separator."""
    if not isinstance(prefix, str):
        raise TypeError("prefix must be a string")

    if not isinstance(addr_hash, (bytes, bytearray)):
        raise TypeError("addr_hash must be binary bytes")

    if kind not in (SCRIPT_TYPE, PUBKEY_TYPE, SCRIPT_TYPE_TOKEN, PUBKEY_TYPE_TOKEN):
        raise ValueError(f"unrecognised address type {kind}")

    payload = _pack_addr_data(kind, addr_hash)
    checksum = _create_checksum(prefix, payload)
    return "".join([_CHARSET[d] for d in (payload + checksum)])


def encode_full(prefix, kind, addr_hash):
    """Encode a full cashaddr address, with prefix and separator."""
    return ":".join([prefix, encode(prefix, kind, addr_hash)])