#!/usr/bin/env python3
# BlockStore: a helper class that keeps a map of blocks and implements
#             helper functions for responding to getheaders and getdata,
#             and for constructing a getheaders message
#

from io import BytesIO
import dbm.ndbm
import sys

from .environment import network, Network

if network() == Network.BCH:
    from .bch.nodemessages import (
        MsgTx,
        CBlock,
        MsgHeaders,
        CBlockHeader,
        MsgBlock,
        CTransaction,
    )
elif network() == Network.NEX:
    from .nex.nodemessages import (
        MsgTx,
        CBlock,
        MsgHeaders,
        CBlockHeader,
        MsgBlock,
        CTransaction,
    )
else:
    raise NotImplementedError()


class BlockStore:
    def __init__(self, datadir):
        self.block_db = dbm.ndbm.open(datadir + "/blocks", "c")
        self.current_block = 0
        self.headers_map = {}

    def close(self):
        self.block_db.close()

    def get(self, blockhash):
        serialized_block = None
        try:
            serialized_block = self.block_db[repr(blockhash)]
        except KeyError:
            return None
        f = BytesIO(serialized_block)
        ret = CBlock()
        ret.deserialize(f)
        ret.calc_hash()
        return ret

    def get_header(self, blockhash):
        try:
            return self.headers_map[blockhash]
        except KeyError:
            return None

    # Note: this pulls full blocks out of the database just to retrieve
    # the headers -- perhaps we could keep a separate data structure
    # to avoid this overhead.
    def headers_for(self, locator, hash_stop, current_tip=None):
        if current_tip is None:
            current_tip = self.current_block
        current_block_header = self.get_header(current_tip)
        if current_block_header is None:
            return None

        response = MsgHeaders()
        headers_list = [current_block_header]
        maxheaders = 2000
        while headers_list[0].gethash() not in locator.vHave:
            prev_block_hash = headers_list[0].hashPrevBlock
            prev_block_header = self.get_header(prev_block_hash)
            if prev_block_header is not None:
                headers_list.insert(0, prev_block_header)
            else:
                break
        headers_list = headers_list[:maxheaders]  # truncate if we have too many
        hash_list = [x.gethash() for x in headers_list]
        index = len(headers_list)
        if hash_stop in hash_list:
            index = hash_list.index(hash_stop) + 1
        response.headers = headers_list[:index]
        return response

    def add_block(self, block):
        block.calc_hash()
        try:
            self.block_db[repr(block.gethash())] = bytes(block.serialize())
        except TypeError as e:
            print("Unexpected error: ", sys.exc_info()[0], e.args)
        self.current_block = block.gethash()
        self.headers_map[block.gethash()] = CBlockHeader(block)

    def add_header(self, header):
        self.headers_map[header.gethash()] = header

    def get_blocks(self, inv):
        responses = []
        for i in inv:
            if i.type == 2:  # MSG_BLOCK
                block = self.get(i.hash)
                if block is not None:
                    responses.append(MsgBlock(block))
        return responses


class TxStore:
    def __init__(self, datadir):
        self.tx_db = dbm.ndbm.open(datadir + "/transactions", "c")

    def close(self):
        self.tx_db.close()

    def get(self, txhash):
        serialized_tx = None
        try:
            serialized_tx = self.tx_db[repr(txhash)]
        except KeyError:
            return None
        f = BytesIO(serialized_tx)
        ret = CTransaction()
        ret.deserialize(f)
        ret.calc_idem()
        return ret

    def add_transaction(self, tx):
        idem = tx.calcIdem()
        try:
            self.tx_db[repr(idem)] = bytes(tx.serialize())
        except TypeError as e:
            print("Unexpected error: ", sys.exc_info()[0], e.args)

    def get_transactions(self, inv):
        responses = []
        for i in inv:
            if i.type == 1:  # MSG_TX
                tx = self.get(i.hash)
                if tx is not None:
                    responses.append(MsgTx(tx))
        return responses