#!/usr/bin/env python3
'''
    Onionr - Private P2P Communication

    Multiprocess proof of work
'''
'''
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
'''

import subprocess, os
import multiprocessing, threading, time, json
from multiprocessing import Pipe, Process
from onionrblocks import onionrblockapi
import config, onionrutils, logger, onionrproofs, onionrcrypto as crypto
from onionrutils import bytesconverter
class SubprocessPOW:
    def __init__(self, data, metadata, subproc_count=None):
        '''
            Onionr proof of work using multiple processes
            Accepts block data, block metadata
            if subproc_count is not set, os.cpu_count() is used to determine the number of processes

            Do to Python GIL multiprocessing or use of external libraries is necessary to accelerate CPU bound tasks
        '''
        # No known benefit to using more processes than there are cores.
        # Note: os.cpu_count perhaps not always accurate
        if subproc_count is None:
            subproc_count = os.cpu_count()
        self.subproc_count = subproc_count
        self.result = ''
        self.shutdown = False
        self.data = data
        self.metadata = metadata

        # dump dict to measure bytes of json metadata. Cannot reuse later because the pow token must be added
        json_metadata = json.dumps(metadata).encode()

        self.data = bytesconverter.str_to_bytes(data)
        # Calculate difficulty. Dumb for now, may use good algorithm in the future.
        self.difficulty = onionrproofs.getDifficultyForNewBlock(bytes(json_metadata + b'\n' + self.data))
        
        logger.info('Computing POW (difficulty: %s)...' % (self.difficulty,))

        self.mainHash = '0' * 64
        self.puzzle = self.mainHash[0:min(self.difficulty, len(self.mainHash))]
        self.shutdown = False
        self.payload = None

    def start(self):
        # Create a new thread for each subprocess
        for x in range(self.subproc_count):
            threading.Thread(target=self._spawn_proc).start()
        # Monitor the processes for a payload, shut them down when its found
        while True:
            if self.payload is None:
                time.sleep(0.1)
            else:
                self.shutdown = True
                return self.payload
    
    def _spawn_proc(self):
        # Create a child proof of work process, wait for data and send shutdown signal when its found
        parent_conn, child_conn = Pipe()
        p = Process(target=self.do_pow, args=(child_conn,))
        p.start()
        p.join()
        payload = None
        try:
            while True:
                data = parent_conn.recv()
                if len(data) >= 1:
                    payload = data
                    break
        except KeyboardInterrupt:
            pass
        finally:
            parent_conn.send('shutdown')
            self.payload = payload

    def do_pow(self, pipe):
        nonce = -10000000 # Start nonce at negative 10 million so that the chosen nonce is likely to be small in length
        nonceStart = nonce
        data = self.data
        metadata = self.metadata
        puzzle = self.puzzle
        difficulty = self.difficulty
        while True:
            # Break if shutdown received
            if pipe.poll() and pipe.recv() == 'shutdown':
                break
            # Load nonce into block metadata
            metadata['pow'] = nonce
            # Serialize metadata, combine with block data
            payload = json.dumps(metadata).encode() + b'\n' + data
            # Check sha3_256 hash of block, compare to puzzle. Send payload if puzzle finished
            token = crypto.hashers.sha3_hash(payload)
            token = bytesconverter.bytes_to_str(token) # ensure token is string
            if puzzle == token[0:difficulty]:
                pipe.send(payload)
                break
            nonce += 1