#! /usr/bin/env python3

import base64, hashlib, os, queue, socket, sys, threading, time

n = 22152137184697602751949152182712806144286735269755991212091578018969267364854563246965161398309375840488156997640625085213985729183180764348127989435514689722834129288342499703533613231239853168289771769536412763137391054558055082146752229593979328251181647873233949602834141648160681711983351545692646009424518816069561938917629523175464947983950548802679152115205735609960641453864298194702935993896839374645356040490091081577992299773430144650589605043643969140352237968606446474316247592579560197155686719175897498255683642121505357781103123719079205647707696181709515150954235402701095586525936356219917713227143

def log(msg):
    time_info = time.strftime("%d.%m.%Y %H:%M:%S")
    print(f"[{time_info}] {msg}")

def load_progress():
    try:
        f = open("progress", "r+b")
    except FileNotFoundError:
        f = open("progress", "w+b")
        return [f, 0, 0, 0]
    length = f.seek(0, os.SEEK_END)
    if length < 256:
        return [f, 0, 0, 0]
    f.seek(0)
    begin = int.from_bytes(f.read(256), "big")
    position = length // 256 - 1
    f.seek(256 * position)
    current_value = int.from_bytes(f.read(256), "big")
    return [f, begin, position, current_value]

def normalize(x):
    if x > n // 2:
        return n - x
    return x

def reduce_milestones(a, e):
    half_length = len(a) // 2
    return [(a[i] * pow(a[half_length+i], e, n)) % n for i in range(half_length+1)]

def finalize(f, revealer_pubkey):
    f.seek(0)
    milestones = [int.from_bytes(f.read(256), "big") for _ in range(2**15 + 1)]
    final_R = normalize(milestones[-1])
    intermediates = []
    for i in range(30, 10, -1):
        H = normalize(milestones[0])
        R = normalize(milestones[-1])
        if i > 15:
            I = normalize(milestones[len(milestones) // 2])
        else:
            I = normalize(pow(H, 2**2**(i-1), n))
            milestones = [milestones[0], I, milestones[-1]]
        hash_data = revealer_pubkey + \
            H.to_bytes(256, "big") + \
            I.to_bytes(256, "big") + \
            R.to_bytes(256, "big") + \
            i.to_bytes(1, "big")
        hash_value = hashlib.sha256(hash_data).digest()
        e = int.from_bytes(hash_value, "big")
        milestones = reduce_milestones(milestones, e)
        intermediates.append(I)
    proof_numbers = [final_R] + intermediates
    return [x.to_bytes(256, "big") for x in proof_numbers]

def send_transaction(revealer_pubkey, proof):
    proof_block = b"".join(proof)
    associated_proof_hash = hashlib.sha256(proof_block).digest()
    transaction = b"\x03" + revealer_pubkey + associated_proof_hash + 84 * b"\0"

    s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
    s.connect(("::1", 62039))

    attempts = 10
    next_attempt = time.time()
    while True:
        timeout = next_attempt - time.time()
        if timeout > 0:
            s.settimeout(timeout)
            try:
                msg = s.recv(4096)
                if len(msg) == 38 and msg[0:2] == b"\0\0" and msg[4] == 15 and msg[6:38] == associated_proof_hash:
                    bitfield = msg[5]
                    for bitpos in range(6):
                        if (bitfield & (1 << bitpos)) != 0:
                            begin = 896 * bitpos
                            end = 896 * (bitpos + 1)
                            reply = b"\0\0\0\0\x10" + bytes([bitpos]) + associated_proof_hash + proof_block[begin:end]
                            s.send(reply)
                elif len(msg) == 5 and msg[0:2] == b"\0\0" and msg[4] == 10:
                    return True
            except TimeoutError:
                pass
        elif attempts == 0:
            return False
        else:
            attempts -= 1
            next_attempt += 1.0
            transaction_request = b"\0\0\0\0\x09" + transaction
            s.send(transaction_request)

def mining_loop(event_queue, revealer_pubkey):
    step_exponent = 2**2**15
    f, begin, position, value = load_progress()
    last_status = position * 100 // 2**15
    proof = None
    sent = False
    def reset():
        nonlocal f, begin, position, value, last_status, proof, sent
        f.seek(0)
        f.truncate()
        if begin != 0:
            f.write(begin.to_bytes(256, "big"))
        position = 0
        last_status = 0
        value = begin
        proof = None
        sent = False
    while True:
        if begin == 0:
            begin = event_queue.get()
            reset()
            continue
        try:
            new_begin = event_queue.get(0)
            if new_begin != begin:
                begin = new_begin
                if begin == 0:
                    log("No task right now")
                else:
                    log("Got a new revealing task")
                reset()
                continue
        except queue.Empty:
            pass
        if position < 2**15:
            value = pow(value, step_exponent, n)
            f.write(value.to_bytes(256, "big"))
            position += 1
            new_status = position * 100 // 2**15
            if new_status != last_status:
                log(f"{new_status} %")
                last_status = new_status
        elif proof is None:
            log("finalizing")
            proof = finalize(f, revealer_pubkey)
        elif not sent:
            log("Sending reveal transaction")
            sent = send_transaction(revealer_pubkey, proof)
            if sent:
                log("Transaction has been received by our node")
            else:
                log("(no response from the node)")
        else:
            new_begin = event_queue.get()
            if new_begin == begin:
                continue
            begin = new_begin
            if begin == 0:
                log("No task right now")
            else:
                log("Got a new revealing task")
            reset()

# communicate with node, send change events
def main():
    if len(sys.argv) < 2:
        print("Usage: ./mine.py <wallet address>", file=sys.stderr)
        exit(1)

    revealer_pubkey = base64.b64decode(sys.argv[1])
    assert len(revealer_pubkey) == 32

    s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
    s.connect(("::1", 62039))
    s.settimeout(1)

    event_queue = None

    while True:
        request = b"\0\0\0\0\x0b" + 32 * b"\0"
        s.send(request)
        try:
            response = s.recv(4096)
            if len(response) == 37 and response[0:5] == b"\0\0\0\0\x0c":
                begin = int.from_bytes(response[5:37], "big")
                if event_queue is None:
                    event_queue = queue.Queue()
                    threading.Thread(target=mining_loop, args=(event_queue, revealer_pubkey)).start()
                    if begin == 0:
                        log("node is available, waiting for a reveal mining task")
                    else:
                        log("node is available, starting reveal mining process")
                event_queue.put(begin)
            time.sleep(1)
        except TimeoutError:
            log("Got no response from the node")

if __name__ == '__main__':
    main()