#!/usr/bin/env python3
"""
test_merkle.py - Unit tests for Merkle tree implementation

Run:
    python -m pytest tests/test_merkle.py -v
    python tests/test_merkle.py  # Standalone
"""

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))

import hashlib
from helix_cdc.receipts.merkle import MerkleTree, StreamingMerkleBuilder


def test_empty_tree():
    """Empty tree should have deterministic root."""
    tree = MerkleTree()
    root = tree.get_root()
    assert root == hashlib.sha256(b"").hexdigest()
    assert len(tree) == 0


def test_single_leaf():
    """Single leaf tree."""
    tree = MerkleTree()
    leaf = "a" * 64
    idx = tree.add_leaf(leaf)

    assert idx == 0
    assert len(tree) == 1

    # Root should be hash of domain-separated leaf
    expected = hashlib.sha256(b"\x00" + bytes.fromhex(leaf)).hexdigest()
    assert tree.get_root() == expected


def test_two_leaves():
    """Two leaf tree should combine correctly."""
    tree = MerkleTree()
    tree.add_leaf("a" * 64)
    tree.add_leaf("b" * 64)

    assert len(tree) == 2
    root = tree.get_root()
    assert len(root) == 64  # SHA256 hex

    # Manually compute expected root
    leaf_a = hashlib.sha256(b"\x00" + bytes.fromhex("a" * 64)).digest()
    leaf_b = hashlib.sha256(b"\x00" + bytes.fromhex("b" * 64)).digest()
    expected = hashlib.sha256(b"\x01" + leaf_a + leaf_b).hexdigest()
    assert root == expected


def test_proof_generation():
    """Proof generation should produce valid proof."""
    tree = MerkleTree()
    leaves = ["a" * 64, "b" * 64, "c" * 64, "d" * 64]
    for leaf in leaves:
        tree.add_leaf(leaf)

    for i in range(4):
        proof = tree.get_proof(i)
        assert proof["leaf_index"] == i
        assert "leaf_hash" in proof
        assert "siblings" in proof
        assert proof["root_hash"] == tree.get_root()


def test_proof_verification():
    """Proof verification should pass for valid proofs."""
    tree = MerkleTree()
    leaves = ["a" * 64, "b" * 64, "c" * 64, "d" * 64, "e" * 64]
    for leaf in leaves:
        tree.add_leaf(leaf)

    root = tree.get_root()

    # Verify each leaf
    for i in range(5):
        proof = tree.get_proof(i)
        is_valid = MerkleTree.verify_proof(
            proof["leaf_hash"],
            proof["siblings"],
            root
        )
        assert is_valid, f"Proof for leaf {i} should be valid"


def test_tamper_detection():
    """Tampered proofs should fail verification."""
    tree = MerkleTree()
    for i in range(4):
        tree.add_leaf(f"{i}" * 64)

    proof = tree.get_proof(1)
    root = tree.get_root()

    # Tamper with leaf hash
    tampered = dict(proof)
    tampered["leaf_hash"] = "f" * 64
    assert not MerkleTree.verify_proof(
        tampered["leaf_hash"],
        tampered["siblings"],
        root
    )

    # Tamper with sibling
    if proof["siblings"]:
        tampered = dict(proof)
        tampered["siblings"] = [{"hash": "f" * 64, "position": "right"}]
        assert not MerkleTree.verify_proof(
            tampered["leaf_hash"],
            tampered["siblings"],
            root
        )


def test_root_caching():
    """Root should be cached and invalidated on new leaf."""
    tree = MerkleTree()
    tree.add_leaf("a" * 64)

    root1 = tree.get_root()
    root2 = tree.get_root()
    assert root1 == root2  # Should use cache

    tree.add_leaf("b" * 64)
    root3 = tree.get_root()
    assert root3 != root1  # Should recompute


def test_odd_number_of_leaves():
    """Odd number of leaves should work (last leaf duplicated)."""
    tree = MerkleTree()
    for i in range(5):
        tree.add_leaf(f"{i}" * 64)

    root = tree.get_root()
    assert len(root) == 64

    # All proofs should verify
    for i in range(5):
        proof = tree.get_proof(i)
        assert MerkleTree.verify_proof(
            proof["leaf_hash"],
            proof["siblings"],
            root
        )


def test_from_leaves():
    """Reconstructing tree from leaves should produce same root."""
    original = MerkleTree()
    leaves = [f"{i}" * 64 for i in range(10)]
    for leaf in leaves:
        original.add_leaf(leaf)

    reconstructed = MerkleTree.from_leaves(leaves)
    assert original.get_root() == reconstructed.get_root()


def test_streaming_builder():
    """Streaming builder should produce same root as regular tree."""
    tree = MerkleTree()
    builder = StreamingMerkleBuilder()

    leaves = [f"{i}" * 64 for i in range(16)]
    for leaf in leaves:
        tree.add_leaf(leaf)
        builder.add(leaf)

    assert tree.get_root() == builder.finalize()


def test_streaming_builder_small():
    """Streaming builder with small number of leaves."""
    for n in range(1, 10):
        tree = MerkleTree()
        builder = StreamingMerkleBuilder()

        leaves = [f"{i}" * 64 for i in range(n)]
        for leaf in leaves:
            tree.add_leaf(leaf)
            builder.add(leaf)

        assert tree.get_root() == builder.finalize(), f"Mismatch for n={n}"


def test_sha3_512_mode():
    """SHA3-512 mode should work."""
    tree = MerkleTree(hash_algo="sha3_512")
    tree.add_leaf("a" * 64)
    tree.add_leaf("b" * 64)

    root = tree.get_root()
    assert len(root) == 128  # SHA3-512 = 64 bytes = 128 hex chars

    proof = tree.get_proof(0)
    assert MerkleTree.verify_proof(
        proof["leaf_hash"],
        proof["siblings"],
        root,
        hash_algo="sha3_512"
    )


def test_proof_sibling_positions():
    """Verify sibling positions are correct."""
    tree = MerkleTree()
    tree.add_leaf("a" * 64)
    tree.add_leaf("b" * 64)

    # Leaf 0 should have leaf 1 as right sibling
    proof0 = tree.get_proof(0)
    assert proof0["siblings"][0]["position"] == "right"

    # Leaf 1 should have leaf 0 as left sibling
    proof1 = tree.get_proof(1)
    assert proof1["siblings"][0]["position"] == "left"


def test_large_tree():
    """Large tree should work efficiently."""
    tree = MerkleTree()
    n = 1000

    for i in range(n):
        tree.add_leaf(hashlib.sha256(f"leaf_{i}".encode()).hexdigest())

    root = tree.get_root()
    assert len(root) == 64

    # Verify random leaf
    proof = tree.get_proof(500)
    assert MerkleTree.verify_proof(
        proof["leaf_hash"],
        proof["siblings"],
        root
    )

    # Proof should have log2(n) ≈ 10 siblings
    assert len(proof["siblings"]) <= 11


# =============================================================================
# STANDALONE TEST RUNNER
# =============================================================================

if __name__ == "__main__":
    import traceback

    tests = [
        test_empty_tree,
        test_single_leaf,
        test_two_leaves,
        test_proof_generation,
        test_proof_verification,
        test_tamper_detection,
        test_root_caching,
        test_odd_number_of_leaves,
        test_from_leaves,
        test_streaming_builder,
        test_streaming_builder_small,
        test_sha3_512_mode,
        test_proof_sibling_positions,
        test_large_tree,
    ]

    passed = 0
    failed = 0

    print("Merkle Tree Tests")
    print("=" * 60)

    for test in tests:
        try:
            test()
            print(f"  ✓ {test.__name__}")
            passed += 1
        except AssertionError as e:
            print(f"  ✗ {test.__name__}: {e}")
            failed += 1
        except Exception as e:
            print(f"  ✗ {test.__name__}: {type(e).__name__}: {e}")
            traceback.print_exc()
            failed += 1

    print()
    print(f"Results: {passed} passed, {failed} failed")
    sys.exit(0 if failed == 0 else 1)
