#!/usr/bin/env python3
"""
merkle.py - Merkle Tree with Incremental Leaf Addition and Proof Generation

Enables random access verification: "Verify step 812 without downloading everything"

Design:
- Incremental: Add leaves one at a time, tree grows automatically
- Proof generation: O(log n) siblings for any leaf
- Verification: Standalone proof verification (no full tree needed)
- Compatible with SHA3-512 from bench_core.py

Usage:
    tree = MerkleTree()

    # Add receipts as they're generated
    tree.add_leaf(receipt_1.receipt_hash())
    tree.add_leaf(receipt_2.receipt_hash())
    ...
    tree.add_leaf(receipt_812.receipt_hash())

    # Get proof for step 812
    proof = tree.get_proof(812)

    # Verify without the full tree
    root = open('published_root.txt').read().strip()
    assert MerkleTree.verify_proof(proof['leaf_hash'], proof['siblings'], root)
"""

from __future__ import annotations
import hashlib
from dataclasses import dataclass
from typing import List, Dict, Any, Optional, Tuple


@dataclass
class MerkleProofEntry:
    """Single sibling in a Merkle proof."""
    hash: str
    position: str  # "left" or "right"


class MerkleTree:
    """
    Incremental Merkle tree with proof generation.

    Uses SHA256 (matching receipt schema) with domain separation:
    - Leaf nodes: H(0x00 || data)
    - Internal nodes: H(0x01 || left || right)
    """

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

    def __init__(self, hash_algo: str = "sha256"):
        """
        Initialize empty Merkle tree.

        Args:
            hash_algo: Hash algorithm ("sha256" or "sha3_512")
        """
        self.hash_algo = hash_algo
        self.leaves: List[str] = []  # Raw leaf data (hex strings)
        self.leaf_hashes: List[bytes] = []  # H(0x00 || leaf)
        self._cached_root: Optional[str] = None

    def _hash(self, data: bytes) -> bytes:
        """Compute hash using configured algorithm."""
        if self.hash_algo == "sha3_512":
            return hashlib.sha3_512(data).digest()
        else:
            return hashlib.sha256(data).digest()

    def _hash_leaf(self, data: str) -> bytes:
        """Hash a leaf with domain separation."""
        return self._hash(self.LEAF_PREFIX + bytes.fromhex(data))

    def _hash_node(self, left: bytes, right: bytes) -> bytes:
        """Hash an internal node with domain separation."""
        return self._hash(self.NODE_PREFIX + left + right)

    def add_leaf(self, leaf_hash: str) -> int:
        """
        Add a leaf to the tree.

        Args:
            leaf_hash: Hex string hash of the leaf data

        Returns:
            Index of the added leaf
        """
        self.leaves.append(leaf_hash)
        self.leaf_hashes.append(self._hash_leaf(leaf_hash))
        self._cached_root = None  # Invalidate cache
        return len(self.leaves) - 1

    def get_root(self) -> str:
        """
        Compute the Merkle root.

        Returns:
            Hex string of the root hash
        """
        if self._cached_root is not None:
            return self._cached_root

        if not self.leaf_hashes:
            # Empty tree
            self._cached_root = self._hash(b"").hex()
            return self._cached_root

        # Build tree bottom-up
        nodes = list(self.leaf_hashes)

        while len(nodes) > 1:
            next_level = []
            it = iter(nodes)
            for left in it:
                right = next(it, left)  # Duplicate last node if odd
                next_level.append(self._hash_node(left, right))
            nodes = next_level

        self._cached_root = nodes[0].hex()
        return self._cached_root

    def get_proof(self, leaf_index: int) -> Dict[str, Any]:
        """
        Generate Merkle proof for a leaf.

        Args:
            leaf_index: Index of the leaf (0-based)

        Returns:
            {
                "leaf_index": int,
                "leaf_hash": str (hex),
                "siblings": [{"hash": str, "position": "left"|"right"}, ...],
                "root_hash": str (hex)
            }
        """
        if leaf_index < 0 or leaf_index >= len(self.leaves):
            raise ValueError(f"Leaf index {leaf_index} out of range [0, {len(self.leaves)})")

        siblings: List[Dict[str, str]] = []
        nodes = list(self.leaf_hashes)
        idx = leaf_index

        # Build proof while traversing up the tree
        while len(nodes) > 1:
            next_level = []
            it = iter(enumerate(nodes))

            for i, left in it:
                try:
                    j, right = next(it)
                except StopIteration:
                    # Odd number of nodes - duplicate last
                    right = left
                    j = i

                # Check if our target leaf is in this pair
                if i == idx or j == idx:
                    if i == idx:
                        siblings.append({
                            "hash": right.hex(),
                            "position": "right"
                        })
                        idx = len(next_level)
                    else:
                        siblings.append({
                            "hash": left.hex(),
                            "position": "left"
                        })
                        idx = len(next_level)

                next_level.append(self._hash_node(left, right))

            nodes = next_level

        return {
            "leaf_index": leaf_index,
            "leaf_hash": self.leaf_hashes[leaf_index].hex(),
            "siblings": siblings,
            "root_hash": self.get_root()
        }

    @classmethod
    def verify_proof(
        cls,
        leaf_hash: str,
        siblings: List[Dict[str, str]],
        root_hash: str,
        hash_algo: str = "sha256"
    ) -> bool:
        """
        Verify a Merkle proof against a known root.

        Args:
            leaf_hash: Hex string of the leaf hash (already domain-separated)
            siblings: List of {"hash": str, "position": "left"|"right"}
            root_hash: Expected Merkle root (hex string)
            hash_algo: Hash algorithm to use

        Returns:
            True if proof is valid
        """
        if hash_algo == "sha3_512":
            def hash_fn(data: bytes) -> bytes:
                return hashlib.sha3_512(data).digest()
        else:
            def hash_fn(data: bytes) -> bytes:
                return hashlib.sha256(data).digest()

        current = bytes.fromhex(leaf_hash)

        for sibling in siblings:
            sibling_bytes = bytes.fromhex(sibling["hash"])
            if sibling["position"] == "right":
                current = hash_fn(cls.NODE_PREFIX + current + sibling_bytes)
            else:
                current = hash_fn(cls.NODE_PREFIX + sibling_bytes + current)

        return current.hex() == root_hash

    def __len__(self) -> int:
        """Number of leaves in the tree."""
        return len(self.leaves)

    def to_dict(self) -> Dict[str, Any]:
        """Serialize tree state."""
        return {
            "hash_algo": self.hash_algo,
            "leaf_count": len(self.leaves),
            "leaves": self.leaves,
            "root_hash": self.get_root(),
        }

    @classmethod
    def from_leaves(cls, leaves: List[str], hash_algo: str = "sha256") -> "MerkleTree":
        """
        Reconstruct tree from list of leaf hashes.

        Args:
            leaves: List of hex string leaf hashes
            hash_algo: Hash algorithm

        Returns:
            Reconstructed MerkleTree
        """
        tree = cls(hash_algo=hash_algo)
        for leaf in leaves:
            tree.add_leaf(leaf)
        return tree


# =============================================================================
# BATCH / STREAMING HELPERS
# =============================================================================

class StreamingMerkleBuilder:
    """
    Memory-efficient streaming Merkle root computation.

    For very large receipt streams where we can't hold all leaves in memory.
    Only stores O(log n) intermediate hashes.
    """

    def __init__(self, hash_algo: str = "sha256"):
        self.hash_algo = hash_algo
        self.pending: List[bytes] = []  # Stack of pending nodes at each level
        self.count = 0

    def _hash(self, data: bytes) -> bytes:
        if self.hash_algo == "sha3_512":
            return hashlib.sha3_512(data).digest()
        return hashlib.sha256(data).digest()

    def _hash_leaf(self, data: str) -> bytes:
        return self._hash(MerkleTree.LEAF_PREFIX + bytes.fromhex(data))

    def _hash_node(self, left: bytes, right: bytes) -> bytes:
        return self._hash(MerkleTree.NODE_PREFIX + left + right)

    def add(self, leaf_hash: str) -> None:
        """Add a leaf incrementally."""
        node = self._hash_leaf(leaf_hash)
        self.count += 1

        # Propagate up the tree
        level = 0
        while level < len(self.pending):
            pending = self.pending[level]
            if pending is None:
                self.pending[level] = node
                return
            else:
                # Combine with pending node
                node = self._hash_node(pending, node)
                self.pending[level] = None
                level += 1

        # New level needed
        self.pending.append(node)

    def finalize(self) -> str:
        """Compute final root after all leaves added."""
        if not self.pending:
            return self._hash(b"").hex()

        # Collect non-None pending nodes with their levels
        # pending[i] represents a subtree of height i with 2^i leaves
        nodes_with_levels = [
            (level, node) for level, node in enumerate(self.pending)
            if node is not None
        ]

        if not nodes_with_levels:
            return self._hash(b"").hex()

        if len(nodes_with_levels) == 1:
            return nodes_with_levels[0][1].hex()

        # To match the tree's behavior:
        # - Tree duplicates odd nodes at each level
        # - We need to combine pending nodes while maintaining correct level ordering
        #
        # Algorithm: Process from lowest level to highest.
        # A node at level L needs to be "promoted" to level L+1 by duplicating it.
        # Then it can be combined with nodes at level L+1.

        # Start with the lowest level node
        result_level, result = nodes_with_levels[0]

        for next_level, next_node in nodes_with_levels[1:]:
            # Promote result to match next_level by duplicating at each step
            while result_level < next_level:
                # Duplicate: combine with self to go up one level
                result = self._hash_node(result, result)
                result_level += 1

            # Now at same level - combine (next_node comes first as it's the "left" subtree)
            result = self._hash_node(next_node, result)
            result_level += 1

        return result.hex()


# =============================================================================
# DEMO
# =============================================================================

if __name__ == "__main__":
    print("Merkle Tree Demo")
    print("=" * 60)

    # Create tree with sample receipts
    tree = MerkleTree()

    receipts = [
        "a" * 64,  # step_0
        "b" * 64,  # step_1
        "c" * 64,  # step_2
        "d" * 64,  # step_3
        "e" * 64,  # step_4
    ]

    print("\n1. Adding leaves:")
    for i, receipt_hash in enumerate(receipts):
        idx = tree.add_leaf(receipt_hash)
        print(f"   Added leaf {idx}: {receipt_hash[:16]}...")

    root = tree.get_root()
    print(f"\n2. Merkle root: {root[:32]}...")

    # Get and verify proof for leaf 2
    print("\n3. Proof for leaf 2 (step_2):")
    proof = tree.get_proof(2)
    print(f"   Leaf hash: {proof['leaf_hash'][:32]}...")
    print(f"   Siblings ({len(proof['siblings'])}):")
    for s in proof["siblings"]:
        print(f"     - {s['position']}: {s['hash'][:32]}...")
    print(f"   Root: {proof['root_hash'][:32]}...")

    # Verify proof
    is_valid = MerkleTree.verify_proof(
        proof["leaf_hash"],
        proof["siblings"],
        proof["root_hash"]
    )
    print(f"\n4. Proof verification: {'PASS' if is_valid else 'FAIL'}")

    # Verify tamper detection
    print("\n5. Tamper detection test:")
    tampered_proof = dict(proof)
    tampered_proof["leaf_hash"] = "f" * 64  # Wrong hash
    is_valid_tampered = MerkleTree.verify_proof(
        tampered_proof["leaf_hash"],
        tampered_proof["siblings"],
        tampered_proof["root_hash"]
    )
    print(f"   Tampered proof verification: {'PASS (BAD!)' if is_valid_tampered else 'FAIL (CORRECT!)'}")

    # Streaming builder demo
    print("\n6. Streaming Merkle Builder:")
    builder = StreamingMerkleBuilder()
    for receipt_hash in receipts:
        builder.add(receipt_hash)
    streaming_root = builder.finalize()
    print(f"   Streaming root: {streaming_root[:32]}...")
    print(f"   Matches tree:   {streaming_root == root}")
