"""Tests for merkle.py: balanced binary Merkle tree, proofs, authenticated ops."""

import hmac
import inspect

import pytest
from hypothesis import assume, given, settings
from hypothesis import strategies as st

from ads_foundation.merkle import (
    EMPTY_TREE_HASH,
    MerkleNode,
    MerkleProof,
    Side,
    build_merkle,
    generate_proof,
    height,
    insert_auth,
    leaves,
    lookup_auth,
    root_hash,
    size,
    verify_proof,
)
from ads_foundation.core import Auth, Mode, ProofStream, prover_context, verifier_context
from ads_foundation.hash import sha256, shallow_hash


# ---------------------------------------------------------------------------
# Strategies
# ---------------------------------------------------------------------------

int_lists = st.lists(
    st.integers(min_value=-1000, max_value=1000), min_size=1, max_size=50
)

str_lists = st.lists(
    st.text(min_size=1, max_size=10, alphabet=st.characters(whitelist_categories=("L",))),
    min_size=1,
    max_size=30,
)


# ---------------------------------------------------------------------------
# Basic correctness
# ---------------------------------------------------------------------------

class TestBuildMerkle:
    def test_empty(self):
        assert build_merkle([]) is None

    def test_singleton(self):
        tree = build_merkle([42])
        assert tree is not None
        assert tree.data == 42
        assert tree.left is None
        assert tree.right is None

    def test_two_elements(self):
        tree = build_merkle([1, 2])
        assert tree is not None
        assert tree.data is None  # internal node
        assert tree.left is not None and tree.right is not None
        assert tree.left.data == 1
        assert tree.right.data == 2

    def test_frozen(self):
        tree = build_merkle([1, 2, 3])
        with pytest.raises(AttributeError):
            tree.hash = b"bad"  # type: ignore[misc]

    @given(int_lists)
    def test_deterministic_root_hash(self, values):
        """Same input -> same root hash."""
        tree1 = build_merkle(values)
        tree2 = build_merkle(values)
        assert root_hash(tree1) == root_hash(tree2)

    @given(int_lists)
    def test_leaves_roundtrip(self, values):
        """leaves(build_merkle(xs)) == xs."""
        tree = build_merkle(values)
        assert leaves(tree) == values

    def test_empty_root_hash(self):
        assert root_hash(None) == EMPTY_TREE_HASH

    @given(int_lists)
    def test_all_hashes_32_bytes(self, values):
        """Every node hash is 32 bytes."""
        tree = build_merkle(values)
        assert len(root_hash(tree)) == 32


class TestSizeHeight:
    @given(int_lists)
    def test_size_matches_input(self, values):
        tree = build_merkle(values)
        assert size(tree) == len(values)

    def test_size_empty(self):
        assert size(None) == 0

    @given(int_lists)
    def test_height_bound(self, values):
        """Height is O(log n) for balanced tree."""
        import math
        tree = build_merkle(values)
        n = len(values)
        # Balanced tree: height <= ceil(log2(n)) + 1
        assert height(tree) <= math.ceil(math.log2(n)) + 1 if n > 1 else height(tree) == 1

    def test_height_empty(self):
        assert height(None) == 0


# ---------------------------------------------------------------------------
# Proof generation and verification
# ---------------------------------------------------------------------------

class TestProofs:
    def test_simple_proof(self):
        data = [10, 20, 30, 40]
        tree = build_merkle(data)
        proof = generate_proof(0, tree, len(data))
        assert verify_proof(10, proof)

    def test_last_element_proof(self):
        data = [10, 20, 30, 40]
        tree = build_merkle(data)
        proof = generate_proof(3, tree, len(data))
        assert verify_proof(40, proof)

    def test_singleton_proof(self):
        tree = build_merkle([99])
        proof = generate_proof(0, tree, 1)
        assert verify_proof(99, proof)
        assert len(proof.path) == 0  # no siblings for singleton

    @given(int_lists)
    def test_completeness(self, values):
        """Every element has a valid proof (completeness)."""
        tree = build_merkle(values)
        for i, v in enumerate(values):
            proof = generate_proof(i, tree, len(values))
            assert verify_proof(v, proof), f"Proof failed for index {i}, value {v}"

    @given(int_lists)
    def test_proof_size_log_bound(self, values):
        """Proof path length <= height - 1 (O(log n))."""
        import math
        tree = build_merkle(values)
        n = len(values)
        for i in range(n):
            proof = generate_proof(i, tree, n)
            assert len(proof.path) <= height(tree) - 1
            if n > 1:
                assert len(proof.path) >= 1

    def test_wrong_value_fails(self):
        data = [10, 20, 30, 40]
        tree = build_merkle(data)
        proof = generate_proof(0, tree, len(data))
        assert not verify_proof(99, proof)  # wrong value

    def test_out_of_range_index(self):
        tree = build_merkle([1, 2, 3])
        with pytest.raises(IndexError):
            generate_proof(3, tree, 3)
        with pytest.raises(IndexError):
            generate_proof(-1, tree, 3)

    def test_empty_tree_proof_raises(self):
        with pytest.raises(ValueError):
            generate_proof(0, MerkleNode(hash=b"\x00" * 32), 0)


# ---------------------------------------------------------------------------
# Tamper detection
# ---------------------------------------------------------------------------

class TestTamperDetection:
    @given(
        st.lists(st.integers(min_value=-100, max_value=100), min_size=2, max_size=30),
        st.integers(min_value=0, max_value=31),
    )
    def test_flip_sibling_byte(self, values, flip_byte):
        """Flipping any byte in a sibling hash breaks verification."""
        tree = build_merkle(values)
        proof = generate_proof(0, tree, len(values))
        if not proof.path:
            return  # singleton, no siblings to tamper
        # Tamper with first sibling
        side, sib_hash = proof.path[0]
        mutated = bytearray(sib_hash)
        mutated[flip_byte] ^= 0xFF
        tampered_path = [(side, bytes(mutated))] + proof.path[1:]
        tampered = MerkleProof(path=tampered_path, root_hash=proof.root_hash)
        assert not verify_proof(values[0], tampered)

    @given(
        st.lists(st.integers(min_value=0, max_value=50), min_size=2, max_size=20),
        st.lists(st.integers(min_value=51, max_value=100), min_size=2, max_size=20),
    )
    def test_proof_non_replayability(self, values1, values2):
        """Proof for tree1 must NOT verify against tree2's root."""
        tree1 = build_merkle(values1)
        tree2 = build_merkle(values2)
        assume(root_hash(tree1) != root_hash(tree2))
        proof1 = generate_proof(0, tree1, len(values1))
        # Transplant root hash from tree2
        transplanted = MerkleProof(path=proof1.path, root_hash=root_hash(tree2))
        assert not verify_proof(values1[0], transplanted)


# ---------------------------------------------------------------------------
# Security properties
# ---------------------------------------------------------------------------

class TestSecurity:
    def test_domain_separation(self):
        """Leaf hash != internal node hash even for crafted input."""
        from ads_foundation.merkle import _LEAF_TAG, _NODE_TAG
        # A leaf with value 1 and an internal node combining two zero hashes
        # should never collide
        leaf = build_merkle([1])
        internal = build_merkle([1, 2])
        assert leaf.hash != internal.hash

    def test_verify_uses_constant_time_compare(self):
        """verify_proof uses hmac.compare_digest (source inspection)."""
        source = inspect.getsource(verify_proof)
        assert "hmac.compare_digest" in source

    def test_all_hashes_32_bytes(self):
        tree = build_merkle([1, 2, 3, 4, 5])
        proof = generate_proof(2, tree, 5)
        assert len(proof.root_hash) == 32
        for _side, sib in proof.path:
            assert len(sib) == 32

    def test_empty_tree_hash_is_constant(self):
        assert EMPTY_TREE_HASH == sha256(b"EmptyTree")
        assert len(EMPTY_TREE_HASH) == 32


# ---------------------------------------------------------------------------
# Authenticated operations
# ---------------------------------------------------------------------------

class TestAuthenticatedOps:
    def test_insert_auth_roundtrip(self):
        """Prover proof accepted by verifier for insert."""
        data = [10, 20, 30]
        tree = build_merkle(data)
        rh = root_hash(tree)

        # Prover
        with prover_context() as p:
            auth_tree = Auth.prover(rh, tree)
            auth_tree = insert_auth(40, auth_tree, p)
        proof = p.get_proof()
        new_hash = auth_tree.digest()

        # Verifier
        with verifier_context(proof) as v:
            auth_tree_v = Auth.verifier(rh)
            auth_tree_v = insert_auth(40, auth_tree_v, v)
        assert auth_tree_v.digest() == new_hash

    def test_multiple_inserts(self):
        """Chain of authenticated inserts."""
        tree = build_merkle([1])
        rh = root_hash(tree)

        with prover_context() as p:
            at = Auth.prover(rh, tree)
            at = insert_auth(2, at, p)
            at = insert_auth(3, at, p)
            at = insert_auth(4, at, p)
        proof = p.get_proof()
        final_hash = at.digest()

        with verifier_context(proof) as v:
            at_v = Auth.verifier(rh)
            at_v = insert_auth(2, at_v, v)
            at_v = insert_auth(3, at_v, v)
            at_v = insert_auth(4, at_v, v)
        assert at_v.digest() == final_hash

    @given(int_lists, st.integers(min_value=-1000, max_value=1000))
    @settings(max_examples=50)
    def test_insert_auth_hypothesis(self, values, x):
        """Prover/verifier roundtrip for arbitrary inserts."""
        tree = build_merkle(values)
        rh = root_hash(tree)

        with prover_context() as p:
            at = Auth.prover(rh, tree)
            at = insert_auth(x, at, p)
        proof = p.get_proof()

        with verifier_context(proof) as v:
            at_v = Auth.verifier(rh)
            at_v = insert_auth(x, at_v, v)
        assert at.digest() == at_v.digest()

    @given(int_lists)
    @settings(max_examples=50)
    def test_mode_equivalence(self, values):
        """Prover and verifier agree on hash after insert."""
        tree = build_merkle(values)
        rh = root_hash(tree)
        x = max(values) + 1  # guaranteed new

        with prover_context() as p:
            at = Auth.prover(rh, tree)
            at = insert_auth(x, at, p)
        proof = p.get_proof()

        with verifier_context(proof) as v:
            at_v = Auth.verifier(rh)
            at_v = insert_auth(x, at_v, v)
        assert at.digest() == at_v.digest()


# ---------------------------------------------------------------------------
# String data (multi-type coverage)
# ---------------------------------------------------------------------------

class TestStringData:
    @given(str_lists)
    def test_string_leaves_roundtrip(self, values):
        tree = build_merkle(values)
        assert leaves(tree) == values

    @given(str_lists)
    def test_string_proofs(self, values):
        tree = build_merkle(values)
        for i, v in enumerate(values):
            proof = generate_proof(i, tree, len(values))
            assert verify_proof(v, proof)
