"""Tests for bst.py: BST operations, proofs, authenticated operations."""

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

from ads_foundation.bst import (
    BST,
    Direction,
    MembershipProof,
    generate_proof,
    insert_auth,
    delete_auth,
    lookup_auth,
    verify_proof,
)
from ads_foundation.core import Auth, Mode, ProofStream, auth, prover_context, verifier_context
from ads_foundation.hash import shallow_hash


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def build_tree(values: list[int]) -> BST[int] | None:
    """Build a BST from a list of values."""
    tree: BST[int] | None = None
    for v in values:
        tree = BST.insert(v, tree)
    return tree


# Strategy for non-empty unique-ish integer lists
int_lists = st.lists(
    st.integers(min_value=-1000, max_value=1000), min_size=1, max_size=50
)


# ---------------------------------------------------------------------------
# Plain BST correctness
# ---------------------------------------------------------------------------

class TestBSTBasic:
    def test_insert_lookup(self):
        tree = build_tree([5, 3, 7, 1, 9])
        assert BST.lookup(5, tree) == 5
        assert BST.lookup(3, tree) == 3
        assert BST.lookup(99, tree) is None

    def test_delete(self):
        tree = build_tree([5, 3, 7, 1, 9])
        tree = BST.delete(5, tree)
        assert BST.lookup(5, tree) is None
        assert BST.lookup(3, tree) == 3

    def test_size(self):
        tree = build_tree([5, 3, 7, 1, 9])
        assert BST.size(tree) == 5

    def test_height(self):
        tree = build_tree([5, 3, 7])
        assert BST.height(tree) >= 1

    def test_inorder(self):
        tree = build_tree([5, 3, 7, 1, 9])
        assert BST.inorder(tree) == [1, 3, 5, 7, 9]

    def test_empty_tree(self):
        assert BST.lookup(1, None) is None
        assert BST.size(None) == 0
        assert BST.height(None) == 0
        assert BST.inorder(None) == []
        assert BST.delete(1, None) is None

    def test_frozen(self):
        tree = BST(value=1)
        with pytest.raises(AttributeError):
            tree.value = 2  # type: ignore[misc]


# ---------------------------------------------------------------------------
# Structural invariants (Hypothesis)
# ---------------------------------------------------------------------------

class TestBSTStructural:
    @given(int_lists, st.integers(min_value=-1000, max_value=1000))
    def test_ordering_preserved(self, values, x):
        """P11: inorder is always sorted after insert."""
        tree = build_tree(values)
        new_tree = BST.insert(x, tree)
        assert BST.inorder(new_tree) == sorted(set(values) | {x})

    @given(int_lists)
    def test_delete_lookup(self, values):
        """P12: delete then lookup returns None."""
        tree = build_tree(values)
        target = values[0]
        assert BST.lookup(target, BST.delete(target, tree)) is None

    @given(int_lists, st.integers(min_value=-1000, max_value=1000))
    def test_insert_lookup_consistency(self, values, x):
        """P5: insert then lookup finds the value."""
        tree = build_tree(values)
        new_tree = BST.insert(x, tree)
        assert BST.lookup(x, new_tree) == x

    @given(int_lists, st.integers(min_value=-1000, max_value=1000))
    def test_hash_changes_on_insert(self, values, x):
        """P13: hash changes when a new element is inserted."""
        tree = build_tree(values)
        assume(x not in values)
        new_tree = BST.insert(x, tree)
        assert shallow_hash(tree) != shallow_hash(new_tree)


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

class TestProofs:
    def test_simple_proof(self):
        tree = build_tree([5, 3, 7, 1, 9])
        root_hash = shallow_hash(tree)
        proof = generate_proof(5, tree)
        assert verify_proof(5, root_hash, proof)

    def test_leaf_proof(self):
        tree = build_tree([5, 3, 7, 1, 9])
        root_hash = shallow_hash(tree)
        proof = generate_proof(1, tree)
        assert verify_proof(1, root_hash, proof)

    @given(int_lists)
    def test_completeness(self, values):
        """P2: every element in the tree has a valid proof."""
        tree = build_tree(values)
        root_hash = shallow_hash(tree)
        for v in set(values):
            proof = generate_proof(v, tree)
            assert verify_proof(v, root_hash, proof), f"Proof failed for {v}"

    @given(
        st.lists(st.integers(min_value=0, max_value=100), min_size=1, max_size=30),
        st.integers(min_value=101, max_value=200),
    )
    def test_soundness(self, values, non_member):
        """P3: non-members fail verification."""
        tree = build_tree(values)
        root_hash = shallow_hash(tree)
        proof = generate_proof(non_member, tree)
        assert not verify_proof(non_member, root_hash, proof)

    @given(int_lists)
    def test_proof_size_bound(self, values):
        """M2: proof path length <= tree height."""
        tree = build_tree(values)
        target = values[0]
        proof = generate_proof(target, tree)
        assert len(proof.path) <= BST.height(tree)


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

class TestTamperDetection:
    @given(
        st.lists(st.integers(min_value=-100, max_value=100), min_size=3, max_size=30),
        st.integers(min_value=0, max_value=31),
    )
    def test_flip_proof_byte(self, values, flip_byte):
        """P6: flipping any sibling hash byte breaks verification."""
        tree = build_tree(values)
        root_hash = shallow_hash(tree)
        target = values[0]
        proof = generate_proof(target, tree)
        if not proof.path:
            # Tamper with left_hash instead
            mutated = bytearray(proof.left_hash)
            mutated[flip_byte] ^= 0xFF
            tampered = MembershipProof(True, bytes(mutated), proof.right_hash, proof.path)
            assert not verify_proof(target, root_hash, tampered)
            return
        direction, sibling_hash, value_hash = proof.path[0]
        mutated = bytearray(sibling_hash)
        mutated[flip_byte] ^= 0xFF
        tampered_path = [(direction, bytes(mutated), value_hash)] + proof.path[1:]
        tampered = MembershipProof(True, proof.left_hash, proof.right_hash, tampered_path)
        assert not verify_proof(target, root_hash, 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):
        """P8: proof for tree1 must NOT verify against tree2's root."""
        tree1 = build_tree(values1)
        tree2 = build_tree(values2)
        assume(shallow_hash(tree1) != shallow_hash(tree2))
        target = values1[0]
        proof = generate_proof(target, tree1)
        assert not verify_proof(target, shallow_hash(tree2), proof)


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

class TestAuthenticatedOps:
    def test_insert_auth_roundtrip(self):
        """P1: prover proof is accepted by verifier for insert."""
        tree = build_tree([5, 3, 7])
        root_hash = shallow_hash(tree)

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

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

    def test_delete_auth_roundtrip(self):
        tree = build_tree([5, 3, 7])
        root_hash = shallow_hash(tree)

        with prover_context() as p:
            auth_tree = Auth.prover(root_hash, tree)
            auth_tree = delete_auth(3, auth_tree, p)
        proof = p.get_proof()
        new_hash = auth_tree.digest()

        with verifier_context(proof) as v:
            auth_tree_v = Auth.verifier(root_hash)
            auth_tree_v = delete_auth(3, auth_tree_v, v)
        assert auth_tree_v.digest() == new_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_tree(values)
        root_hash = shallow_hash(tree)

        with prover_context() as p:
            auth_tree = Auth.prover(root_hash, tree)
            auth_tree = insert_auth(x, auth_tree, p)
        proof = p.get_proof()
        new_hash = auth_tree.digest()

        with verifier_context(proof) as v:
            auth_tree_v = Auth.verifier(root_hash)
            auth_tree_v = insert_auth(x, auth_tree_v, v)
        assert auth_tree_v.digest() == new_hash

    def test_multiple_inserts(self):
        """Chain of authenticated inserts."""
        tree = build_tree([5])
        root_hash = shallow_hash(tree)

        with prover_context() as p:
            at = Auth.prover(root_hash, tree)
            at = insert_auth(3, at, p)
            at = insert_auth(7, at, p)
            at = insert_auth(1, at, p)
        proof = p.get_proof()
        final_hash = at.digest()

        with verifier_context(proof) as v:
            at_v = Auth.verifier(root_hash)
            at_v = insert_auth(3, at_v, v)
            at_v = insert_auth(7, at_v, v)
            at_v = insert_auth(1, at_v, v)
        assert at_v.digest() == final_hash


# ---------------------------------------------------------------------------
# Metamorphic relations
# ---------------------------------------------------------------------------

class TestMetamorphic:
    @given(int_lists)
    def test_mode_equivalence_insert(self, values):
        """M1: prover and verifier agree on hash after insert."""
        tree = build_tree(values)
        root_hash = shallow_hash(tree)
        x = values[0] + 10000  # likely new

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

        with verifier_context(proof) as v:
            at_v = Auth.verifier(root_hash)
            at_v = insert_auth(x, at_v, v)

        assert at.digest() == at_v.digest()
