"""Append-only Merkle transparency log — the ABOM Notary's tamper-evidence core.

This is the structurally-neutral primitive a regulated buyer's security team
actually checks. Unlike the linked hash chain in ``audit.py`` (where an operator
who controls the store can rebuild the whole chain from genesis), a Merkle
transparency log lets *any* party verify two independent things without trusting
the operator:

  * **inclusion proof**  — "record R is in the log at position i, under signed
    tree head with root H" — an O(log n) audit path.
  * **consistency proof** — "the log of size m is a prefix of the log of size n;
    nothing already published was edited or removed" — append-only enforcement.

The construction follows RFC 6962 (Certificate Transparency): domain-separated
leaf (0x00) and node (0x01) hashes so a leaf can never be forged as an internal
node. The tree is recomputed from the leaf list; for the MVP that is in-process
and O(n) per head, which is fine at demo / design-partner scale. Phase-2 swaps
the leaf store for an external witness + periodic anchoring without changing
these proof formats.

Pure stdlib (hashlib only) so it is unit-testable with no DB and no keys.
"""
from __future__ import annotations

import base64
import hashlib
import os
from dataclasses import dataclass, field
from pathlib import Path

LEAF_PREFIX = b"\x00"
NODE_PREFIX = b"\x01"


def _sha256(data: bytes) -> bytes:
    return hashlib.sha256(data).digest()


def hash_leaf(data: bytes) -> bytes:
    """RFC 6962 leaf hash: SHA256(0x00 || data)."""
    return _sha256(LEAF_PREFIX + data)


def hash_children(left: bytes, right: bytes) -> bytes:
    """RFC 6962 node hash: SHA256(0x01 || left || right)."""
    return _sha256(NODE_PREFIX + left + right)


def _largest_power_of_two_below(n: int) -> int:
    """Largest k that is a power of two and strictly less than n (n >= 2)."""
    k = 1
    while k * 2 < n:
        k *= 2
    return k


def merkle_root(leaves: list[bytes]) -> bytes:
    """Merkle Tree Hash (MTH) of a list of already-hashed leaves (RFC 6962 §2.1)."""
    n = len(leaves)
    if n == 0:
        # MTH of the empty tree is SHA256() of the empty string.
        return _sha256(b"")
    if n == 1:
        return leaves[0]
    k = _largest_power_of_two_below(n)
    return hash_children(merkle_root(leaves[:k]), merkle_root(leaves[k:]))


def inclusion_proof(leaves: list[bytes], index: int) -> list[bytes]:
    """Audit path proving ``leaves[index]`` is in the tree (RFC 6962 §2.1.1)."""
    n = len(leaves)
    if not 0 <= index < n:
        raise IndexError(f"index {index} out of range for tree of size {n}")
    if n == 1:
        return []
    k = _largest_power_of_two_below(n)
    if index < k:
        return inclusion_proof(leaves[:k], index) + [merkle_root(leaves[k:])]
    return inclusion_proof(leaves[k:], index - k) + [merkle_root(leaves[:k])]


def verify_inclusion(leaf: bytes, index: int, tree_size: int, proof: list[bytes], root: bytes) -> bool:
    """Verify an audit path: recompute the root from ``leaf`` + ``proof``."""
    if not 0 <= index < tree_size:
        return False
    fn, sn = index, tree_size - 1
    node = leaf
    for sibling in proof:
        if fn % 2 == 1 or fn == sn:
            node = hash_children(sibling, node)
            while fn % 2 == 0 and fn != 0:
                fn >>= 1
                sn >>= 1
        else:
            node = hash_children(node, sibling)
        fn >>= 1
        sn >>= 1
    return fn == 0 and node == root


def consistency_proof(leaves: list[bytes], m: int, n: int) -> list[bytes]:
    """Proof that the size-``m`` log is a prefix of the size-``n`` log (RFC 6962 §2.1.2)."""
    if not 0 < m <= n <= len(leaves):
        raise ValueError(f"require 0 < m <= n <= len(leaves); got m={m}, n={n}, len={len(leaves)}")
    return _subproof(m, leaves[:n], True)


def _subproof(m: int, leaves: list[bytes], b: bool) -> list[bytes]:
    n = len(leaves)
    if m == n:
        return [] if b else [merkle_root(leaves)]
    k = _largest_power_of_two_below(n)
    if m <= k:
        return _subproof(m, leaves[:k], b) + [merkle_root(leaves[k:])]
    return _subproof(m - k, leaves[k:], False) + [merkle_root(leaves[:k])]


def verify_consistency(m: int, n: int, proof: list[bytes], old_root: bytes, new_root: bytes) -> bool:
    """Verify a consistency proof between roots of sizes ``m`` and ``n``."""
    if m < 1 or m > n:
        return False
    if m == n:
        return not proof and old_root == new_root
    if not proof:
        return False

    # RFC 6962 §2.1.2 verification algorithm.
    if m & (m - 1) == 0:  # m is a power of two: prepend old_root as the seed.
        proof = [old_root] + proof

    fn, sn = m - 1, n - 1
    while fn % 2 == 1:
        fn >>= 1
        sn >>= 1

    fr = sr = proof[0]
    for c in proof[1:]:
        if sn == 0:
            return False
        if fn % 2 == 1 or fn == sn:
            fr = hash_children(c, fr)
            sr = hash_children(c, sr)
            while fn % 2 == 0 and fn != 0:
                fn >>= 1
                sn >>= 1
        else:
            sr = hash_children(sr, c)
        fn >>= 1
        sn >>= 1

    return sn == 0 and fr == old_root and sr == new_root


@dataclass
class TreeHead:
    """A (size, root) pair — what a signer attests to. Root is hex for transport."""
    tree_size: int
    root_hash: str

    def as_dict(self) -> dict:
        return {"tree_size": self.tree_size, "root_hash": self.root_hash}


@dataclass
class MerkleLog:
    """In-process append-only transparency log over raw entry bytes.

    Append returns the leaf index; ``head()`` returns the current signed-able
    tree head; ``prove_inclusion`` / ``prove_consistency`` emit hex proofs that
    any third party can verify with the ``verify_*`` helpers above.
    """
    leaves: list[bytes] = field(default_factory=list)   # leaf hashes
    entries: list[bytes] = field(default_factory=list)  # raw entry bytes (for re-verification)

    def append(self, entry: bytes) -> int:
        self.entries.append(entry)
        self.leaves.append(hash_leaf(entry))
        return len(self.leaves) - 1

    @property
    def size(self) -> int:
        return len(self.leaves)

    def head(self) -> TreeHead:
        return TreeHead(self.size, merkle_root(self.leaves).hex())

    def prove_inclusion(self, index: int) -> dict:
        proof = inclusion_proof(self.leaves, index)
        return {
            "leaf_index": index,
            "tree_size": self.size,
            "audit_path": [h.hex() for h in proof],
            "root_hash": merkle_root(self.leaves).hex(),
        }

    def prove_consistency(self, old_size: int) -> dict:
        proof = consistency_proof(self.leaves, old_size, self.size)
        return {
            "old_size": old_size,
            "new_size": self.size,
            "proof": [h.hex() for h in proof],
            "new_root": merkle_root(self.leaves).hex(),
        }


def verify_inclusion_hex(entry: bytes, proof: dict) -> bool:
    """Verify a hex inclusion proof (as emitted by ``MerkleLog.prove_inclusion``)."""
    return verify_inclusion(
        leaf=hash_leaf(entry),
        index=proof["leaf_index"],
        tree_size=proof["tree_size"],
        proof=[bytes.fromhex(h) for h in proof["audit_path"]],
        root=bytes.fromhex(proof["root_hash"]),
    )


def verify_consistency_hex(old_root_hex: str, proof: dict) -> bool:
    """Verify a hex consistency proof (as emitted by ``MerkleLog.prove_consistency``)."""
    return verify_consistency(
        m=proof["old_size"],
        n=proof["new_size"],
        proof=[bytes.fromhex(h) for h in proof["proof"]],
        old_root=bytes.fromhex(old_root_hex),
        new_root=bytes.fromhex(proof["new_root"]),
    )


class PersistentMerkleLog(MerkleLog):
    """A ``MerkleLog`` backed by an append-only file, so the transparency log
    survives process exit.

    The audit's finding: the in-memory ``MerkleLog`` evaporated on exit, which
    made the RFC 6962 consistency machinery dead code — nothing ever compared two
    *persisted* roots. This subclass makes the entries durable so a later process
    can reload the exact log and prove append-only growth across the restart.

    Storage: one base64-encoded raw entry per line (entries are arbitrary bytes).
    The entries are the source of truth; leaf hashes and roots are recomputed on
    load, so a tampered line changes the reloaded root — which is the point.

    Durability: each append writes one line and fsyncs, so a crash loses at most
    the entry that was in flight (same discipline as the provenance sink).
    """

    def __init__(self, path: str | Path):
        super().__init__()
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._load()

    def _load(self) -> None:
        if not self.path.exists():
            return
        for line in self.path.read_text().splitlines():
            line = line.strip()
            if not line:
                continue
            entry = base64.b64decode(line)
            # rebuild in-memory state via the parent (no re-persist)
            super().append(entry)

    def append(self, entry: bytes) -> int:
        index = super().append(entry)
        with open(self.path, "a", encoding="utf-8") as fh:
            fh.write(base64.b64encode(entry).decode() + "\n")
            fh.flush()
            os.fsync(fh.fileno())
        return index
