"""Blocks, hash pointers, proof-of-work, and the rule two nodes use to decide
which of two perfectly valid chains is history.

Everything is counted in SHA-256 evaluations. `HASHES` is the only clock in
this toy: there is no wall clock and no randomness anywhere in this file, and
the nonce search starts at 0 and increments, so every number the module
produces is byte-identical on every run.

Deliberately absent: a Merkle root in the block header. Committing a whole set
of transactions to one 32-byte field is the neighbouring toy (`merkle-tree/`),
and putting it here would blur two mechanisms together. `data` is opaque bytes,
so nothing distracts from the accounting.
"""

import hashlib

# ---------------------------------------------------------------- primitives

HASHES = [0]          # every SHA-256 in this module goes through H(), so the
                      # counter *is* the work, and no number on the page is a
                      # wall-clock measurement.


def H(b):
    HASHES[0] += 1
    return hashlib.sha256(b).digest()


def counter_reset():
    HASHES[0] = 0


def counted():
    return HASHES[0]


# ---------------------------------------------------------------- the block

class Block:
    """A header and nothing else. `data` stands in for the body.

    `bits` -- the number of leading zero bits this block *claims* to have
    mined against -- lives inside the hashed header, exactly like Bitcoin's
    `nBits`. It is self-declared: the block says how hard it was, and the
    proof-of-work check takes it at its word. That is the whole toy.
    """

    __slots__ = ("index", "prev", "data", "bits", "nonce")

    def __init__(self, index, prev, data, bits, nonce=0):
        self.index = index
        self.prev = prev          # bytes: hash of the previous block's header
        self.data = data          # bytes: the payload
        self.bits = bits          # int: leading zero bits this block claims
        self.nonce = nonce

    def header(self):
        return b"|".join([
            str(self.index).encode(),
            self.prev,
            self.data,
            str(self.bits).encode(),
            str(self.nonce).encode(),
        ])

    def hash(self):
        return H(self.header())

    def __repr__(self):
        return f"Block({self.index}, bits={self.bits}, nonce={self.nonce}, {self.data!r})"


GENESIS_PREV = b"\x00" * 32


def meets(digest, bits):
    """Does this digest clear a target of `bits` leading zero bits?"""
    return int.from_bytes(digest, "big") < (1 << (256 - bits))


def work(bits):
    """Expected hashes to find one block at this difficulty: 2**bits."""
    return 1 << bits


# ---------------------------------------------------------------- mining

def mine(index, prev, data, bits, start=0):
    """Search nonces from `start` upward. Returns (block, hashes spent).

    Starting at 0 rather than at a random nonce is what makes every mining
    cost on the commentary page reproducible: the search is a pure function of
    the header.
    """
    before = counted()
    b = Block(index, prev, data, bits, start)
    while not meets(b.hash(), bits):
        b.nonce += 1
    return b, counted() - before


def build_chain(datas, bits, prev=GENESIS_PREV):
    """Mine a chain of len(datas) blocks. Returns (chain, hashes spent)."""
    before = counted()
    chain = []
    for i, d in enumerate(datas):
        b, _ = mine(i, prev, d, bits)
        chain.append(b)
        prev = b.hash()
    return chain, counted() - before


# ---------------------------------------------------------------- validation

def validate(chain, expected_bits=None):
    """Is this a well-formed chain? Returns (bool, reason).

    Checks, in order:
      - block 0 points at GENESIS_PREV
      - every block's `prev` equals the previous block's hash (hash pointers)
      - every block's hash meets the difficulty *the block itself declares*
      - indices are consecutive

    Note what the third check does NOT do: compare the declared difficulty
    against anything the network expects. `expected_bits`, if given, adds that
    -- and it is a counterfactual, not the default. Leaving it off is what the
    demo's act 3 walks through.
    """
    if not chain:
        return False, "empty"
    prev = GENESIS_PREV
    for i, b in enumerate(chain):
        if b.index != i:
            return False, f"block {i} index {b.index}"
        if b.prev != prev:
            return False, f"block {i} broken hash pointer"
        h = b.hash()
        if not meets(h, b.bits):
            return False, f"block {i} fails its own difficulty {b.bits}"
        if expected_bits is not None and b.bits != expected_bits:
            return False, f"block {i} declares {b.bits}, expected {expected_bits}"
        prev = h
    return True, "valid"


# ---------------------------------------------------------------- fork choice

def total_work(chain):
    """Cumulative expected hashes behind this chain -- Bitcoin's nChainWork."""
    return sum(work(b.bits) for b in chain)


def choose_longest(a, b):
    """The rule everyone repeats: 'longest chain wins'. Ties go to `a`."""
    return a if len(a) >= len(b) else b


def choose_heaviest(a, b):
    """The rule Bitcoin implements: most cumulative work wins. Ties go to `a`."""
    return a if total_work(a) >= total_work(b) else b


# ---------------------------------------------------------------- tampering

def tamper_and_remine(chain, k, new_data, bits=None):
    """Rewrite block k's data and re-mine blocks k..end so the chain relinks.

    `bits=None` keeps each block's original difficulty (an honest rewrite,
    priced at the real rate). Passing `bits` lets the rewriter declare its own
    difficulty, which is the attack in act 3.

    Returns (new chain, hashes spent).
    """
    before = counted()
    out = list(chain[:k])
    prev = chain[k - 1].hash() if k else GENESIS_PREV
    for i in range(k, len(chain)):
        d = new_data if i == k else chain[i].data
        b, _ = mine(i, prev, d, chain[i].bits if bits is None else bits)
        out.append(b)
        prev = b.hash()
    return out, counted() - before