#!/usr/bin/env python3

# Copyright (C) 2017-2020 The btclib developers
#
# This file is part of btclib. It is subject to the license terms in the
# LICENSE file found in the top-level directory of this distribution.
#
# No part of btclib including this file, may be copied, modified, propagated,
# or distributed except according to the terms contained in the LICENSE file.

"""Bitcoin Script.

https://en.bitcoin.it/wiki/Script

Scripts are represented by Iterable[Token] (usually List[Token]),
where Token = Union[int, str, bytes]:

* int [-1, 16] are shorcuts for 'OP_1NEGATE', 'OP_0' - 'OP_16'
* str are for opcodes (e.g. 'OP_HASH160') or hexstring data
* bytes are for data (but integers are often casted to int)
"""

OP_CODES = {
    # Constants
    'OP_0'                  : b'\x00',
    'OP_FALSE'              : b'\x00',
    'OP_PUSHDATA1'          : b'\x4c',
    'OP_PUSHDATA2'          : b'\x4d',
    'OP_PUSHDATA4'          : b'\x4e',
    'OP_1NEGATE'            : b'\x4f',
    'OP_1'                  : b'\x51',
    'OP_TRUE'               : b'\x51',
    'OP_2'                  : b'\x52',
    'OP_3'                  : b'\x53',
    'OP_4'                  : b'\x54',
    'OP_5'                  : b'\x55',
    'OP_6'                  : b'\x56',
    'OP_7'                  : b'\x57',
    'OP_8'                  : b'\x58',
    'OP_9'                  : b'\x59',
    'OP_10'                 : b'\x5a',
    'OP_11'                 : b'\x5b',
    'OP_12'                 : b'\x5c',
    'OP_13'                 : b'\x5d',
    'OP_14'                 : b'\x5e',
    'OP_15'                 : b'\x5f',
    'OP_16'                 : b'\x60',

    # Flow control
    'OP_NOP'                : b'\x61',
    'OP_IF'                 : b'\x63',
    'OP_NOTIF'              : b'\x64',
    'OP_ELSE'               : b'\x67',
    'OP_ENDIF'              : b'\x68',
    'OP_VERIFY'             : b'\x69',
    'OP_RETURN'             : b'\x6a',

    # Stack
    'OP_TOALTSTACK'         : b'\x6b',
    'OP_FROMALTSTACK'       : b'\x6c',
    'OP_IFDUP'              : b'\x73',
    'OP_DEPTH'              : b'\x74',
    'OP_DROP'               : b'\x75',
    'OP_DUP'                : b'\x76',
    'OP_NIP'                : b'\x77',
    'OP_OVER'               : b'\x78',
    'OP_PICK'               : b'\x79',
    'OP_ROLL'               : b'\x7a',
    'OP_ROT'                : b'\x7b',
    'OP_SWAP'               : b'\x7c',
    'OP_TUCK'               : b'\x7d',
    'OP_2DROP'              : b'\x6d',
    'OP_2DUP'               : b'\x6e',
    'OP_3DUP'               : b'\x6f',
    'OP_2OVER'              : b'\x70',
    'OP_2ROT'               : b'\x71',
    'OP_2SWAP'              : b'\x72',

    # Splice
    'OP_SIZE'               : b'\x82',

    # Bitwise logic
    'OP_EQUAL'              : b'\x87',
    'OP_EQUALVERIFY'        : b'\x88',

    # Arithmetic
    'OP_1ADD'               : b'\x8b',
    'OP_1SUB'               : b'\x8c',
    'OP_NEGATE'             : b'\x8f',
    'OP_ABS'                : b'\x90',
    'OP_NOT'                : b'\x91',
    'OP_0NOTEQUAL'          : b'\x92',
    'OP_ADD'                : b'\x93',
    'OP_SUB'                : b'\x94',
    'OP_BOOLAND'            : b'\x9a',
    'OP_BOOLOR'             : b'\x9b',
    'OP_NUMEQUAL'           : b'\x9c',
    'OP_NUMEQUALVERIFY'     : b'\x9d',
    'OP_NUMNOTEQUAL'        : b'\x9e',
    'OP_LESSTHAN'           : b'\x9f',
    'OP_GREATERTHAN'        : b'\xa0',
    'OP_LESSTHANOREQUAL'    : b'\xa1',
    'OP_GREATERTHANOREQUAL' : b'\xa2',
    'OP_MIN'                : b'\xa3',
    'OP_MAX'                : b'\xa4',
    'OP_WITHIN'             : b'\xa5',

    # Crypto
    'OP_RIPEMD160'          : b'\xa6',
    'OP_SHA1'               : b'\xa7',
    'OP_SHA256'             : b'\xa8',
    'OP_HASH160'            : b'\xa9',
    'OP_HASH256'            : b'\xaa',
    'OP_CODESEPARATOR'      : b'\xab',
    'OP_CHECKSIG'           : b'\xac',
    'OP_CHECKSIGVERIFY'     : b'\xad',
    'OP_CHECKMULTISIG'      : b'\xae',
    'OP_CHECKMULTISIGVERIFY': b'\xaf',

    # Locktime
    'OP_NOP2'               : b'\xb1',
    'OP_CHECKLOCKTIMEVERIFY': b'\xb1',
    'OP_NOP3'               : b'\xb2',
    'OP_CHECKSEQUENCEVERIFY': b'\xb2',

    # Reserved words
    'OP_NOP1'               : b'\xb0',
    'OP_NOP4'               : b'\xb3',
    'OP_NOP5'               : b'\xb4',
    'OP_NOP6'               : b'\xb5',
    'OP_NOP7'               : b'\xb6',
    'OP_NOP8'               : b'\xb7',
    'OP_NOP9'               : b'\xb8',
    'OP_NOP10'              : b'\xb9'
}

OP_CODE_NAMES = {
      0: 'OP_0',
     76: 'OP_PUSHDATA1',
     77: 'OP_PUSHDATA2',
     78: 'OP_PUSHDATA4',
     79: 'OP_1NEGATE',
     81: 'OP_1',
     82: 'OP_2',
     83: 'OP_3',
     84: 'OP_4',
     85: 'OP_5',
     86: 'OP_6',
     87: 'OP_7',
     88: 'OP_8',
     89: 'OP_9',
     90: 'OP_10',
     91: 'OP_11',
     92: 'OP_12',
     93: 'OP_13',
     94: 'OP_14',
     95: 'OP_15',
     96: 'OP_16',
     97: 'OP_NOP',
     99: 'OP_IF',
    100: 'OP_NOTIF',
    103: 'OP_ELSE',
    104: 'OP_ENDIF',
    105: 'OP_VERIFY',
    106: 'OP_RETURN',
    107: 'OP_TOALTSTACK',
    108: 'OP_FROMALTSTACK',
    109: 'OP_2DROP',
    110: 'OP_2DUP',
    111: 'OP_3DUP',
    112: 'OP_2OVER',
    113: 'OP_2ROT',
    114: 'OP_2SWAP',
    115: 'OP_IFDUP',
    116: 'OP_DEPTH',
    117: 'OP_DROP',
    118: 'OP_DUP',
    119: 'OP_NIP',
    120: 'OP_OVER',
    121: 'OP_PICK',
    122: 'OP_ROLL',
    123: 'OP_ROT',
    124: 'OP_SWAP',
    125: 'OP_TUCK',
    130: 'OP_SIZE',
    135: 'OP_EQUAL',
    136: 'OP_EQUALVERIFY',
    139: 'OP_1ADD',
    140: 'OP_1SUB',
    143: 'OP_NEGATE',
    144: 'OP_ABS',
    145: 'OP_NOT',
    146: 'OP_0NOTEQUAL',
    147: 'OP_ADD',
    148: 'OP_SUB',
    154: 'OP_BOOLAND',
    155: 'OP_BOOLOR',
    156: 'OP_NUMEQUAL',
    157: 'OP_NUMEQUALVERIFY',
    158: 'OP_NUMNOTEQUAL',
    159: 'OP_LESSTHAN',
    160: 'OP_GREATERTHAN',
    161: 'OP_LESSTHANOREQUAL',
    162: 'OP_GREATERTHANOREQUAL',
    163: 'OP_MIN',
    164: 'OP_MAX',
    165: 'OP_WITHIN',
    166: 'OP_RIPEMD160',
    167: 'OP_SHA1',
    168: 'OP_SHA256',
    169: 'OP_HASH160',
    170: 'OP_HASH256',
    171: 'OP_CODESEPARATOR',
    172: 'OP_CHECKSIG',
    173: 'OP_CHECKSIGVERIFY',
    174: 'OP_CHECKMULTISIG',
    175: 'OP_CHECKMULTISIGVERIFY',
    176: 'OP_NOP1',
    177: 'OP_CHECKLOCKTIMEVERIFY',
    178: 'OP_CHECKSEQUENCEVERIFY',
    179: 'OP_NOP4',
    180: 'OP_NOP5',
    181: 'OP_NOP6',
    182: 'OP_NOP7',
    183: 'OP_NOP8',
    184: 'OP_NOP9',
    185: 'OP_NOP10',
}

from io import BytesIO
from typing import BinaryIO, Iterable, List, Union

from . import varint
from .utils import Octets, bytes_from_hexstring

# the integers [0-16] are shorcuts for 'OP_0'-'OP_16'
# the integer -1 is a shorcut for 'OP_1NEGATE'
# other integers are bytes encoded (require push operation) 
# str are for opcodes (e.g. 'OP_HASH160') or hexstring data to be pushed
# bytes is for data to be pushed
Token = Union[int, str, bytes]
# TODO: use Script as input and serialize Iterable[Token] to bytes on the fly
Script = Union[bytes, Iterable[Token]]


def _op_pushdata(data: bytes) -> bytes:
    """Convert to canonical push: OP_PUSHDATA (if needed) | length | data.

    According to standardness rules (BIP-62) the
    minimum possible PUSHDATA operator must be used.
    Byte vectors on the stack are not allowed to be more than 520 bytes long.
    """

    r = b''
    length = len(data)
    if length < 75:     # 1-byte-length
        r += length.to_bytes(1, byteorder='little')
    elif length < 256:  # OP_PUSHDATA1 | 1-byte-length
        r += OP_CODES['OP_PUSHDATA1']
        r += length.to_bytes(1, byteorder='little')
    elif length < 521:  # OP_PUSHDATA2 | 2-byte-length
        r += OP_CODES['OP_PUSHDATA2']
        r += length.to_bytes(2, byteorder='little')
    else:
        # because of the 520 bytes limit
        # there is no need to use OP_PUSHDATA4
        # r += OP_CODES['OP_PUSHDATA4']
        # r += length.to_bytes(4, byteorder='little')
        raise ValueError(f"Script: Cannot push {length} bytes on the stack")
    r += data
    return r


def encode(script: Iterable[Token]) -> bytes:
    r = b''
    for token in script:
        if isinstance(token, int):
            # Short 1-byte opcodes exist
            # to push numbers in [-1, 16]
            if token >= 0 and token <= 16:
                r += OP_CODES['OP_' + str(token)]
            elif token == -1:
                r += OP_CODES['OP_1NEGATE']
            # Pushing any other number requires an
            # explicit push operation of its bytes encoding
            # FIXME: negative numbers?
            else:
                nbytes = (token.bit_length() + 7) // 8
                data = token.to_bytes(nbytes, byteorder='little')
                r += _op_pushdata(data)
        elif isinstance(token, str):
            token = token.strip()
            token = token.upper()
            if token in OP_CODES:
                r += OP_CODES[token]
            else:
                try:
                   token = bytes.fromhex(token)
                except Exception:
                    raise ValueError(f"Script: invalid {token!r} opcode")
                r += _op_pushdata(token)
        elif isinstance(token, bytes):
                r += _op_pushdata(token)
        else:
            raise ValueError(f"Script: unmanaged {type(token)} token type")
    return r


def decode(script: Octets) -> List[Token]:

    script = bytes_from_hexstring(script)

    # initialize the result list
    r: List[Union[str, int, bytes]] = []

    length = len(script)
    s = BytesIO(script)  # TODO: avoid unnecessary streamification
    counter = 0
    while counter < length:
        # get one byte and convert it to an integer
        i = s.read(1)[0]
        counter += 1
        if i == 0:
            # numeric value 0 (OP_0)
            r.append(i)
        elif i == 79:
            # numeric value -1 (OP_1NEGATE)
            r.append(-1)
        elif i > 80 and i < 97:
            # numeric values 1-16 (OP_1-OP_16)
            r.append(i-80)
        elif i < 76:
            # 1-byte-data-length | data
            data = s.read(i)
            counter += i
            n = int.from_bytes(data, byteorder='little')
            if n <= 0xffffffff:
                # probably a number, it won't hurt to consider it a number
                r.append(n)
            else:
                r.append(data.hex())
        elif i == 76:
            # OP_PUSHDATA1 | 1-byte-data-length | data
            data_length = int.from_bytes(s.read(1), byteorder='little')
            data = s.read(data_length)
            counter += 1 + data_length
            r.append(data.hex())
        elif i == 77:
            # OP_PUSHDATA2 | 2-byte-data-length | data
            data_length = int.from_bytes(s.read(2), byteorder='little')
            data = s.read(data_length)
            counter += 2 + data_length
            r.append(data.hex())
        elif i == 78:
            # OP_PUSHDATA4 | 4-byte-data-length | data
            data_length = int.from_bytes(s.read(4), byteorder='little')
            data = s.read(data_length)
            counter += 4 + data_length
            r.append(data.hex())
        else:
            # OP_CODE
            r.append(OP_CODE_NAMES[i])

    return r


def serialize(script: Iterable[Token]) -> bytes:
    r = encode(script)
    # prepend length as varint
    length  = len(r)
    return varint.encode(length) + r


def parse(stream: Union[BinaryIO, Octets]) -> List[Token]:

    stream = bytes_from_hexstring(stream)

    if isinstance(stream, bytes):
        stream = BytesIO(stream)

    length = varint.decode(stream)
    script = stream.read(length)
    return decode(script)
