from dataclasses import dataclass
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey, Ed25519PublicKey
from cryptography.exceptions import InvalidSignature
import hashlib
import time

@dataclass
class Transaction:
    id: int
    sender: bytes
    receiver: bytes
    amount: int
    transaction_fee: int
    signature: bytes

    def is_valid(self):
        sender_pubkey = Ed25519PublicKey.from_public_bytes(self.sender)
        msg = self.id.to_bytes(4, "big") + \
            self.sender + \
            self.receiver + \
            self.amount.to_bytes(8, "big") + \
            self.transaction_fee(8, "big")
        try:
            sender_pubkey.verify(self.signature, msg)
        except InvalidSignature:
            return False
        return self.amount >= 1
    def is_valid_after_block(self, block):
        if (self.sender, self.id) in block.used_transaction_ids:
            return False
        balance = block.balances.get(self.sender)
        if balance is None:
            return False
        return balance >= self.amount + self.transaction_fee
    def get_transaction_raw(self):
        return self.id.to_bytes(4, "big") + \
            self.sender + \
            self.receiver + \
            self.amount.to_bytes(8, "big") + \
            self.transaction_fee.to_bytes(8, "big") + \
            self.signature

@dataclass
class Block:
    nonce: int
    timestamp: int
    previous_hash: bytes
    message: bytes
    difficulty_sum: int
    miner_pubkey: bytes
    transaction: Transaction
    balances: dict
    # (sender_pubkey, id) tuples
    used_transaction_ids: set
    valid: bool

    def validate(self, blockchain):
        if self.transaction is not None:
            if not self.transaction.is_valid():
                return False
        if self.previous_hash != 32 * b"\0":
            prev_block = blockchain.get_block(self.previous_hash)
            if not prev_block.valid:
                return False
            if self.timestamp <= prev_block.timestamp:
                return False
            if self.timestamp > time.time():
                return False
            if not self.transaction.is_valid_after_block(prev_block):
                return False
        else:
            if self.transaction is not None:
                return False
            # check for the correct miner pubkey - which will become public at launch day
            h = hashlib.sha256(self.miner_pubkey).hexdigest()
            if h != "88023d392db35f2d3936abd0532003ae0a38b4d35e4d123a0fa28c568c7e3e2f":
                return False
        B_1_difficulty_sum, B_1_timestamp = self.get_difficulty_info(1, blockchain)
        B_10_difficulty_sum, B_10_timestamp = self.get_difficulty_info(10, blockchain)
        D = B_1_difficulty_sum - B_10_difficulty_sum
        T = self.timestamp - B_10_timestamp
        calculated_difficulty = D * 3000 // 9 // T
        block_difficulty = max(calculated_difficulty, 2**28)
        if B_1_difficulty_sum + block_difficulty != self.difficulty_sum:
            return False
        block_raw = self.get_block_raw()
        block_hash = hashlib.sha256(block_raw).digest()
        self.valid = int.from_bytes(block_hash, "big") * block_difficulty < 2**256
        if self.valid:
            self.calculate_balances(prev_block)
            self.calculate_used_transaction_ids(prev_block)
        return self.valid
    def calculate_balances(self, prev_block):
        if prev_block is None:
            self.balances = {
                self.miner_pubkey: 100,
            }
            return
        balances = prev_block.balances.copy()
        balances.setdefault(self.miner_pubkey, 0)
        balances[miner_pubkey] += 100
        t = self.transaction
        if t is not None:
            balances[miner_pubkey] += t.transaction_fee
            balances[t.sender] -= (t.amount + t.transaction_fee)
            balances[t.receiver] += t.amount
        self.balances = balances
    def calculate_used_transaction_ids(self, prev_block):
        if prev_block is None:
            self.used_transaction_ids = set()
            return
        used_transaction_ids = prev_block.used_transaction_ids.copy()
        t = self.transaction
        if t is not None:
            used_transaction_ids.add((t.sender, t.id))
        self.used_transaction_ids = used_transaction_ids
    def get_difficulty_info(self, steps, blockchain):
        if steps == 0:
            return self.difficulty_sum, self.timestamp
        if self.previous_hash == 32 * b"\0":
            difficulty_sum = 2**29 - steps * 2**28
            timestamp = self.timestamp - steps * 300
            return difficulty_sum, timestamp
        else:
            previous_block = blockchain.get_block(self.previous_hash)
            return previous_block.get_difficulty_info(steps-1, blockchain)
    def get_block_raw(self):
        if self.transaction is None:
            transaction = 148 * b"\0"
        else:
            transaction = self.transaction.get_transaction_raw()
        return self.nonce.to_bytes(8, "big") + \
            self.timestamp.to_bytes(8, "big") + \
            self.previous_hash + \
            self.message + \
            self.difficulty_sum.to_bytes(32, "big") + \
            self.miner_pubkey + \
            transaction

class Blockchain:
    def __init__(self):
        # maps block hashes to block instances
        self.__block_map = {}
    def get_block(self, hash):
        return self.__block_map.get(hash)