#! /usr/bin/env python3

import base64, socket, sys, time
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey

def format_address(raw_address):
    return base64.b64encode(raw_address).decode()

def format_amount(amount, show_plus = True):
    color_prefix = ""
    color_suffix = ""
    sign = "+ " if show_plus else ""
    if amount < 0:
        color_prefix = "\x1b[31m"
        color_suffix = "\x1b[0m"
        sign = "- "
        amount = -amount
    coins = amount / 100
    return f"{color_prefix}{sign}{coins:.02f} cc{color_suffix}"

def write_transaction(timestamp, message, amount):
    formatted_time = time.strftime("%d.%m.%Y %H:%M:%S", time.localtime(timestamp))
    print(f"{formatted_time}  {message:<44}  {format_amount(amount):>14}")

def show_balance(public_key):
    public_key_raw = public_key.public_bytes_raw()
    with open("blockchain", "rb") as f:
        total_amount = 0
        while True:
            block = f.read(292)
            if len(block) != 292:
                break
            miner = block[180:212]
            timestamp = int.from_bytes(block[244:252], "big")
            if block[0:148] != 148 * b"\0":
                sender = block[4:36]
                receiver = block[36:68]
                amount = int.from_bytes(block[68:76], "big")
                fee = int.from_bytes(block[76:84], "big")
                if sender == public_key_raw:
                    write_transaction(timestamp, format_address(receiver), - amount - fee)
                    total_amount -= (amount + fee)
                if receiver == public_key_raw:
                    write_transaction(timestamp, format_address(sender), amount)
                    total_amount += amount
            else:
                fee = 0
            if miner == public_key_raw:
                write_transaction(timestamp, "mining reward", 100 + fee)
                total_amount += 100 + fee
        print(81 * "\u2500")
        amount_string = f"\U0001f955 \x1b[1;37m{format_amount(total_amount, False)}\x1b[0m"
        print(21 * " " + f"\x1b[1;37mYour balance:\x1b[0m{amount_string:>57}")

def parse_amount(amount):
    amount = amount.replace(",", ".")
    parts = amount.split(".")
    if len(parts) == 1:
        return int(parts[0]) * 100
    elif len(parts) == 2:
        if len(parts[1]) > 2:
            raise Exception(f"Invalid amount: {amount}")
        coins = int(parts[0])
        cents = int(parts[1])
        return coins * 100 + cents
    raise Exception(f"Invalid amount: {amount}")

def parse_amount_checked(amount):
    amount = parse_amount(amount)
    if amount < 0:
        raise Exception("amount must not be negative")
    if amount >= 2**64:
        raise Exception("amount is too large")
    return amount

def find_free_id(public_key_raw):
    try:
        with open("blockchain", "rb") as f:
            used_ids = set()
            while True:
                block = f.read(292)
                if len(block) != 292:
                    break
                if block[0:148] != 148 * b"\0":
                    transaction_id = int.from_bytes(block[0:4], "big")
                    sender = block[4:36]
                    if sender == public_key_raw:
                        used_ids.add(transaction_id)
            for possible_id in range(0, 2**32):
                if possible_id not in used_ids:
                    return possible_id
            raise Exception("No transaction id available")
    except FileNotFoundError:
        return 0

def send_payment(private_key, target, amount, fee):
    target_raw = base64.b64decode(target.encode())
    if len(target_raw) != 32:
        raise Exception(f"Invalid target address: {target}")
    amount = parse_amount_checked(amount)
    if amount == 0:
        raise Exception("Amount must not be zero")
    fee = parse_amount_checked(fee)
    public_key_raw = private_key.public_key().public_bytes_raw()
    transaction_id = find_free_id(public_key_raw)
    transaction_prefix = transaction_id.to_bytes(4, "big") + \
        public_key_raw + \
        target_raw + \
        amount.to_bytes(8, "big") + \
        fee.to_bytes(8, "big")
    signature = private_key.sign(transaction_prefix)
    transaction = transaction_prefix + signature
    request = b"\0\0\0\0\x09" + transaction

    s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
    try:
        s.connect(("::1", 62039))
        s.settimeout(1)
        for _ in range(10):
            s.send(request)
            try:
                response = s.recv(1024)
                if response == b"\0\0\0\0\x0a":
                    print("Payment has been transmitted to the node")
                    return
            except TimeoutError:
                pass
        print("node did not respond")
    except ConnectionRefusedError:
        print("node is not running")

def usage_info():
    print("Usage:", file=sys.stderr)
    print("  ./wallet.py # see your past transactions and balance", file=sys.stderr)
    print("  ./wallet.py pay <target> <amount> <fee> # send carrotcoins to someone else", file=sys.stderr)

def main():
    try:
        with open("wallet-key", "rb") as f:
            private_key_raw = f.read()
        private_key = Ed25519PrivateKey.from_private_bytes(private_key_raw)
    except FileNotFoundError:
        private_key = Ed25519PrivateKey.generate()
        private_key_raw = private_key.private_bytes_raw()
        with open("wallet-key", "wb") as f:
            f.write(private_key_raw)

    public_key = private_key.public_key()
    wallet_addr = format_address(public_key.public_bytes_raw())
    print(f"Your wallet address: {wallet_addr}\n")

    if len(sys.argv) == 1:
        show_balance(public_key)
    elif len(sys.argv) == 5 and sys.argv[1] == "pay":
        send_payment(private_key, *sys.argv[2:5])
    else:
        usage_info()
        exit(1)

if __name__ == '__main__':
    main()