#!/usr/bin/env python2.7

import os
import json
import string
import hashlib
import argparse
from sys import argv, stdin
from decimal import Decimal
from functools import partial
from subprocess import Popen, PIPE
from collections import OrderedDict
from binascii import (a2b_hex, b2a_hex)

from crypto import (
    bytes_to_int,
    rand_int,
    secp256k1,
    ec_pub_key,
    ecdsa_sign,
)

# Magic numbers
DEFAULT_FEE = int(1e4)
COIN = int(1e8)
VERSION_P2PKH = chr(0)
#VERSION_P2PKH = chr(111) # testnet
VERSION_SECRET = chr(128)
#VERSION_SECRET = chr(239) # testnet
SIGHASH_ALL = 1
TX_VERSION = 1
INPUT_SEQ_FINAL = 0xffffffff

EC_CURVE = secp256k1

##################################################
# Wallet

class InsufficientFunds(Exception):
    pass

# Use OrderedDicts with integer + ASCII values for easy JSON conversion

def CurvePoint(x, y):
    return OrderedDict([('x', x), ('y', y)])

def Signature(r, s):
    return OrderedDict([('r', r), ('s', s)])

#def Input(txid, index, value, pub_key=None, sig=None):
def Input(txid, index, value, sig_script=None):
    return OrderedDict([
        ('txid', txid),
        ('index', index),
        ('value', value),       # Not included in protocol serialization
        #('pub_key', pub_key),   # We only issue standard signature scripts
        #('signature', sig),
        ('sig_script', sig_script or []),
    ])

def Output(value, address, validate=True):
    if validate:
        a2b_base58check(address)
    return OrderedDict([
        ('value', value),
        ('address', address)    # We only issue and accept standard P2PKH scripts
    ])

def Transaction(inputs, outputs):
    return OrderedDict([
        ('inputs', inputs),
        ('outputs', outputs)
    ])

# For inbound transactions (those paying one or more of our addresses but not
# issued by us), we don't need to store inputs and other outputs.

def PartialOutput(index, value, address):
    return OrderedDict([
        ('index', index),
        ('value', value),
        ('address', address)
    ])

def PartialTransaction(txid, partial_outputs):
    return OrderedDict([
        ('txid', txid),
        ('partial_outputs', partial_outputs),
    ])

class Wallet(object):

    def __init__(self, wallet_path=None):
        if wallet_path is None:
            wallet_path = os.environ.get('WALLET_DIR', '.')
        self._path = wallet_path
        self._keys = None
        self._keys_dirty = False
        self._transactions = None
        self._transactions_dirty = False

    def __enter__(self):
        self._keys = []
        path = os.path.join(self._path, 'keys')
        if os.path.exists(path):
            with open(path, 'rb') as f:
                for line in f:
                    self._keys.append(line.strip())
        self._transactions = []
        path = os.path.join(self._path, 'transactions')
        if os.path.exists(path):
            with open(path, 'rb') as f:
                self._transactions = json.load(f, object_pairs_hook=OrderedDict)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            return False
        if self._keys_dirty:
            keys = ''.join(key + '\n' for key in self._keys)
            path = os.path.join(self._path, 'keys')
            print 'Writing %s' % path
            atomic_overwrite(path, keys)
        if self._transactions_dirty:
            transactions = json.dumps(self._transactions, indent=2)
            path = os.path.join(self._path, 'transactions')
            print 'Writing %s' % path
            atomic_overwrite(path, transactions)

    def keys(self):
        return tuple(self._keys)

    def send_to(self, outputs, fee):
        total = sum(o['value'] for o in outputs) + fee
        inputs = select_inputs(total, self.list_unspent())
        input_total = sum(i['value'] for i in inputs)
        assert input_total >= total
        change = input_total - total
        outputs = list(outputs)
        if change > 0:
            outputs.append(Output(change, self.get_change_address()))
        shuffle(inputs)
        shuffle(outputs)
        tx = Transaction(inputs, outputs)
        self.sign_transaction(tx)
        self._transactions.append(tx)
        self._transactions_dirty = True
        print json.dumps(tx, indent=4)
        print b2a_hex(encode_transaction(tx))

    def list_unspent(self):
        # XXX stub
        return [Input('abcdef', 0, 100000),
                Input('012345', 1, 200000)]

    def get_change_address(self):
        # XXX stub
        return b2a_base58check('\x00changeaddr')

    def sign_transaction(self, tx):
        # Per http://en.bitcoin.it/wiki/OP_CHECKSIG
        def strip_script(i):
            # Return copy of input i with empty signature script (we have to
            # sign all inputs without yet knowing their signatures)
            i = i.copy()
            i['sig_script'] = []
            return i
        inputs = tx['inputs']
        for i, input_ in enumerate(inputs):
            prev_out = self.get_output(input_['txid'], input_['index'])
            stripped_inputs = map(strip_script, inputs)
            stripped_inputs[i]['sig_script'] = output_script(prev_out)
            # ^ I have no idea why we do this. It prevents having a common hash
            # to sign with each key. Satoshi braindamage?
            stripped_tx = Transaction(stripped_inputs, tx['outputs'])
            msg = encode_transaction(stripped_tx) + pack_u32le(SIGHASH_ALL)
            priv_key = self.get_key_by_address(prev_out['address'])
            input_['sig_script'] = input_script(
                    ec_pub_key(EC_CURVE, priv_key),
                    ecdsa_sign(EC_CURVE, sha256d(msg), priv_key, os.urandom))

    def get_output(self, txid, index):
        # XXX slow
        txid = a2b_hex(reversed(txid)) # XXX little endian??
        for tx in self._transactions:
            if sha256d(encode_transaction(tx)) == txid:
                return tx['outputs'][index]
        raise KeyError(txid)

    def get_key_by_address(self, address):
        # XXX slow
        for priv_key in self._keys:
            if bitcoin_address(
                    encode_curve_point(
                        ec_pub_key(priv_key))) == address:
                return priv_key

def main(progname):
    if progname == 'point-addr':
        return point_to_addr()
    if progname == 'key-wif':
        return keys_to_wif()

    parser = argparse.ArgumentParser()
    parser.add_argument('address')
    parser.add_argument('amount')
    parser.add_argument('-b', '--btc', action='store_true',
                        help='specify amounts in BTC (1e8 satoshi)')
    args = parser.parse_args()
    if args.btc:
        args.amount = int(Decimal(args.amount)*COIN)
    else:
        args.amount = int(args.amount)
    with Wallet() as wallet:
        wallet.send_to([Output(args.amount, args.address)], DEFAULT_FEE)

def select_inputs(value, inputs):
    # Choose the smallest input large enough to cover value
    inputs = sorted(inputs, key=lambda i: i['value'])
    for i in inputs:
        if i['value'] >= value:
            return [i]
    # If none large enough, merge from largest to smallest
    selected = []
    total = 0
    for i in reversed(inputs):
        selected.append(i)
        total += i['value']
        if total >= value:
            return selected
    raise InsufficientFunds


##################################################
# Utility functions

_sha256 = hashlib.sha256

def sha256(x):
    return _sha256(x).digest()

def sha256d(x):
    return _sha256(_sha256(x).digest()).digest()

def ripemd160(x):
    # hashlib not guaranteed to have ripemd160 (requires openssl)
    proc = Popen(os.path.join(os.path.dirname(__file__) or '.',
                              'rmd160pipe'),
                 stdin=PIPE, stdout=PIPE)
    out, err = proc.communicate(x)
    if proc.returncode != 0:
        raise RuntimeError('rmd160pipe', err)
    return out

def shuffle(l, rand_bytes=os.urandom):
    "Knuth/Fisher-Yates shuffle in place"
    for i in range(len(l) - 1, 0, -1):
        j = rand_int(i + 1, rand_bytes)
        l[i], l[j] = l[j], l[i]

def atomic_overwrite(path, contents):
    new_path = path + '.new'
    with open(new_path, 'wb') as f:
        f.write(contents)
        os.fsync(f.fileno())
    os.rename(new_path, path)


##################################################
# Base58
# http://en.bitcoin.it/wiki/Base58Check_encoding

base58_alphabet = (string.digits + string.uppercase + string.lowercase).translate(None, '0OIl')
base58_inverse = [None]*256
def init_base58_inverse():
    for index, character in enumerate(base58_alphabet):
        base58_inverse[ord(character)] = index
init_base58_inverse()

def b2a_base58check(data):
    data += sha256d(data)[:4]

    leading_zeros = 0
    for b in data:
        if b != '\x00':
            break
        leading_zeros += 1

    data_num = bytes_to_int(data)

    digits = []
    while data_num:
        data_num, digit = divmod(data_num, 58)
        digits.append(digit)
    digits.extend([0] * leading_zeros)

    return ''.join(base58_alphabet[digit] for digit in reversed(digits))

class Base58Error(ValueError):
    pass

class BadDigit(Base58Error):
    pass

class BadChecksum(Base58Error):
    pass

def a2b_base58(data):
    digits = [base58_inverse[ord(b)] for b in data]
    if None in digits:
        raise BadDigit

    leading_zeros = 0
    for digit in digits:
        if digit != 0:
            break
        leading_zeros += 1

    data_num = 0
    for digit in digits:
        data_num = 58*data_num + digit

    data_bytes = []
    while data_num:
        data_num, byte = divmod(data_num, 256)
        data_bytes.append(byte)
    data_bytes.extend([0] * leading_zeros)

    return ''.join(chr(b) for b in reversed(data_bytes))

def a2b_base58check(data):
    data = a2b_base58(data)
    payload = data[:-4]
    check = data[-4:]
    if check != sha256d(payload)[:4]:
        raise BadChecksum
    return payload

##################################################
# Bitcoin protocol
# http://en.bitcoin.it/wiki/Protocol_documentation

def int_to_fixed_bytes_le(length, i):
    "Convert integer i to fixed-length, zero-padded, little-endian byte field"
    assert i > 0 and i.bit_length() <= 8*length
    return ''.join(chr(i & (0xff << (8*byte))) for byte in range(length))

def int_to_fixed_bytes(length, i):
    "Convert integer i to fixed-length, zero-padded, big-endian byte field"
    return reversed(int_to_fixed_bytes_le(length, i))

pack_u16le = partial(int_to_fixed_bytes_le, 2)
pack_u32le = partial(int_to_fixed_bytes_le, 4)
pack_u64le = partial(int_to_fixed_bytes_le, 8)
pack_u256be = partial(int_to_fixed_bytes, 32)

def encode_var_int(i):
    if i < 0xfd:
        return chr(i)
    elif i < 2**16:
        return '\xfd' + pack_u16le(i)
    elif i < 2**32:
        return '\xfe' + pack_u32le(i)
    else:
        return '\xff' + pack_u64le(i)

def encode_var_str(s):
    return encode_var_int(len(s)) + s

OPCODE_MAP = {
    'OP_PUSHDATA1':     0x4c,
    'OP_PUSHDATA2':     0x4d,
    'OP_PUSHDATA4':     0x4e,
    'OP_RETURN':        0x6a,
    'OP_DUP':           0x76,
    'OP_EQUALVERIFY':   0x88,
    'OP_HASH160':       0xa9,
    'OP_CHECKSIG':      0xac,
}

class ScriptData(str):
    pass

def encode_script_op(op):
    if isinstance(op, ScriptData):
        nbytes = len(op)
        # Because one compact integer encoding just isn't enough...
        if nbytes < OPCODE_MAP['OP_PUSHDATA1']:
            return chr(nbytes) + op
        elif nbytes < 2**8:
            return OPCODE_MAP['OP_PUSHDATA1'] + chr(nbytes) + op
        elif nbytes < 2**16:
            return OPCODE_MAP['OP_PUSHDATA2'] + pack_s16le(nbytes) + op
        elif nbytes < 2**32:
            return OPCODE_MAP['OP_PUSHDATA4'] + pack_s32le(nbytes) + op
        else:
            raise OverflowError
    else:
        return OPCODE_MAP[op]

def encode_script(ops):
    return ''.join(encode_script_op(op) for op in ops)

def encode_curve_point(point, compress=False):
    x, y = point
    if compress:
        if y & 1 == 0:
            sign = '\x02'
        else:
            sign = '\x03'
        return sign + pack_u256be(x)
    else:
        return '\x04' + pack_u256be(x) + pack_u256be(y)

def encode_sig(sig):
    r, s = sig
    # XXX stub DER encoding
    der = pack_u256be(r) + pack_u256be(s)
    return der + chr(SIGHASH_ALL)

def input_script(pub_key, signature, compress=False):
    "Build standard pay-to-pubkey-hash signature script for an input"
    return (ScriptData(encode_sig(signature)),
            ScriptData(encode_curve_point(pub_key, compress)))

def output_script(output):
    "Build standard pay-to-pubkey-hash script for an output"
    return ('OP_DUP', 'OP_HASH160',
            ScriptData(a2b_base58check(output['address'])),
            'OP_EQUALVERIFY', 'OP_CHECKSIG')

def encode_transaction(tx):
    fields = [pack_u32le(TX_VERSION)] # Transaction format version
    inputs = tx['inputs']
    fields.append(encode_var_int(len(inputs)))
    for i in inputs:
        fields.extend((
            i['txid'],
            pack_u32le(i['index']),
            encode_var_str(encode_script(i['sig_script'])),
            pack_u32le(INPUT_SEQ_FINAL)
        ))
    outputs = tx['outputs']
    fields.append(encode_var_int(len(outputs)))
    for o in outputs:
        fields.extend((
            pack_u64le(o['value']),
            encode_var_str(encode_script(output_script(o)))
        ))
    fields.append(pack_u32le(0)) # We don't use lock_time
    return ''.join(fields)

def bitcoin_address(encoded_point):
    return b2a_base58check(VERSION_P2PKH +
                           ripemd160(sha256(encoded_point)))

def point_to_addr():
    print bitcoin_address(stdin.read())

def keys_to_wif():
    for line in stdin:
        bin_key = a2b_hex(line.strip())
        if len(bin_key) != 32:
            raise ValueError("bad key length")
        print b2a_base58check(VERSION_SECRET + bin_key +
                              chr(1)) # for compressed pubkey

if __name__ == '__main__':
    main(os.path.basename(argv[0]))