"""Authenticated Binary Search Tree.

Ports ADS4All's BST.hs: plain BST operations, shallow hashing,
authenticated insert/lookup/delete, and standalone proof generation/verification.
"""

from __future__ import annotations

import hmac
from dataclasses import dataclass
from enum import Enum, auto
from typing import TypeVar

from .core import Auth, Mode, ProofStream
from .hash import combine_hashes, sha256, shallow_hash

T = TypeVar("T")

LEAF_HASH: bytes = sha256(b"Leaf")
"""Sentinel hash for empty trees (Leaf nodes)."""


class Direction(Enum):
    LEFT = auto()
    RIGHT = auto()


# Each proof entry: (direction_taken, sibling_subtree_hash, node_value_hash)
# This gives the verifier enough info to reconstruct the full node hash:
#   combine_hashes(sha256(b"Node"), left_hash, value_hash, right_hash)
ProofEntry = tuple[Direction, bytes, bytes]
ProofPath = list[ProofEntry]


@dataclass(frozen=True, slots=True)
class MembershipProof:
    """Complete proof that a value exists in a BST.

    Includes the found node's child hashes (so the verifier can reconstruct
    the node's hash) plus the path from that node up to the root.
    """

    found: bool
    left_hash: bytes   # left child hash of the found node (LEAF_HASH if None)
    right_hash: bytes  # right child hash of the found node (LEAF_HASH if None)
    path: ProofPath    # path from root down to the found node


@dataclass(frozen=True, slots=True)
class BST[T]:
    """Binary Search Tree node. None represents an empty tree (Leaf)."""

    value: T
    left: BST[T] | None = None
    right: BST[T] | None = None

    def __ads_hash__(self) -> bytes:
        """Shallow hash of BST node. Mirrors Haskell's Shallow instance."""
        left_hash = shallow_hash(self.left) if self.left else LEAF_HASH
        right_hash = shallow_hash(self.right) if self.right else LEAF_HASH
        return combine_hashes(
            sha256(b"Node"),
            left_hash,
            shallow_hash(self.value),
            right_hash,
        )

    @staticmethod
    def insert(x: T, tree: BST[T] | None) -> BST[T]:
        """Functional insert -- returns a new tree."""
        if tree is None:
            return BST(value=x)
        if x < tree.value:  # type: ignore[operator]
            return BST(tree.value, BST.insert(x, tree.left), tree.right)
        elif x > tree.value:  # type: ignore[operator]
            return BST(tree.value, tree.left, BST.insert(x, tree.right))
        else:
            return BST(x, tree.left, tree.right)

    @staticmethod
    def lookup(x: T, tree: BST[T] | None) -> T | None:
        """Search for a value. Returns it if found, None otherwise."""
        if tree is None:
            return None
        if x < tree.value:  # type: ignore[operator]
            return BST.lookup(x, tree.left)
        elif x > tree.value:  # type: ignore[operator]
            return BST.lookup(x, tree.right)
        else:
            return tree.value

    @staticmethod
    def delete(x: T, tree: BST[T] | None) -> BST[T] | None:
        """Functional delete -- returns a new tree."""
        if tree is None:
            return None
        if x < tree.value:  # type: ignore[operator]
            return BST(tree.value, BST.delete(x, tree.left), tree.right)
        elif x > tree.value:  # type: ignore[operator]
            return BST(tree.value, tree.left, BST.delete(x, tree.right))
        else:
            if tree.left is None:
                return tree.right
            if tree.right is None:
                return tree.left
            successor = BST._min_value(tree.right)
            return BST(successor, tree.left, BST.delete(successor, tree.right))

    @staticmethod
    def _min_value(tree: BST[T]) -> T:
        current = tree
        while current.left is not None:
            current = current.left
        return current.value

    @staticmethod
    def size(tree: BST[T] | None) -> int:
        if tree is None:
            return 0
        return 1 + BST.size(tree.left) + BST.size(tree.right)

    @staticmethod
    def height(tree: BST[T] | None) -> int:
        if tree is None:
            return 0
        return 1 + max(BST.height(tree.left), BST.height(tree.right))

    @staticmethod
    def inorder(tree: BST[T] | None) -> list[T]:
        if tree is None:
            return []
        return BST.inorder(tree.left) + [tree.value] + BST.inorder(tree.right)


# ---------------------------------------------------------------------------
# Standalone proof generation and verification (no ProofStream needed)
# ---------------------------------------------------------------------------


def generate_proof(target: T, tree: BST[T] | None) -> MembershipProof:
    """Generate a Merkle-style membership proof for a BST lookup.

    Returns a MembershipProof containing:
    - Whether the value was found
    - The found node's child hashes (for reconstruction)
    - The path from root to the found node with sibling/value hashes
    """
    path: ProofPath = []
    node = tree
    while node is not None:
        value_hash = shallow_hash(node.value)
        if target < node.value:  # type: ignore[operator]
            sibling_hash = shallow_hash(node.right) if node.right else LEAF_HASH
            path.append((Direction.LEFT, sibling_hash, value_hash))
            node = node.left
        elif target > node.value:  # type: ignore[operator]
            sibling_hash = shallow_hash(node.left) if node.left else LEAF_HASH
            path.append((Direction.RIGHT, sibling_hash, value_hash))
            node = node.right
        else:
            # Found: record child hashes of the found node
            left_h = shallow_hash(node.left) if node.left else LEAF_HASH
            right_h = shallow_hash(node.right) if node.right else LEAF_HASH
            return MembershipProof(found=True, left_hash=left_h, right_hash=right_h, path=path)
    # Not found: ended at a leaf
    return MembershipProof(found=False, left_hash=LEAF_HASH, right_hash=LEAF_HASH, path=path)


def verify_proof(value: T, root_hash: bytes, proof: MembershipProof) -> bool:
    """Verify a membership proof by reconstructing the root hash.

    Uses the found node's child hashes + value to build the node hash,
    then walks the proof path from leaf to root, combining with sibling
    hashes at each level. The final hash must equal the trusted root_hash.
    """
    assert len(root_hash) == 32, f"Root hash must be 32 bytes, got {len(root_hash)}"
    if not proof.found:
        return False
    node_tag = sha256(b"Node")
    # Start with the found node's full hash
    current = combine_hashes(
        node_tag, proof.left_hash, shallow_hash(value), proof.right_hash
    )
    for direction, sibling_hash, value_hash in reversed(proof.path):
        assert len(sibling_hash) == 32
        assert len(value_hash) == 32
        if direction == Direction.LEFT:
            # We went left: current is left subtree, sibling is right subtree
            current = combine_hashes(node_tag, current, value_hash, sibling_hash)
        else:
            # We went right: sibling is left subtree, current is right subtree
            current = combine_hashes(node_tag, sibling_hash, value_hash, current)
    return hmac.compare_digest(current, root_hash)


# ---------------------------------------------------------------------------
# Authenticated BST operations (use ProofStream)
# ---------------------------------------------------------------------------


def insert_auth(
    x: T,
    auth_tree: Auth[BST[T] | None],
    ctx: ProofStream,
) -> Auth[BST[T] | None]:
    """Authenticated insert.

    Prover: performs insert, appends NEW root hash to proof, returns new Auth.
    Verifier: consumes proof hash as new root hash.
    """
    if ctx.mode == Mode.PROVER:
        tree = auth_tree.unwrap()
        new_tree = BST.insert(x, tree)
        new_hash = shallow_hash(new_tree)
        ctx.append(new_hash)
        return Auth.prover(new_hash, new_tree)
    else:
        new_hash = ctx.consume()
        assert len(new_hash) == 32
        return Auth.verifier(new_hash)


def delete_auth(
    x: T,
    auth_tree: Auth[BST[T] | None],
    ctx: ProofStream,
) -> Auth[BST[T] | None]:
    """Authenticated delete. Same pattern as insert_auth."""
    if ctx.mode == Mode.PROVER:
        tree = auth_tree.unwrap()
        new_tree = BST.delete(x, tree)
        new_hash = shallow_hash(new_tree) if new_tree else LEAF_HASH
        ctx.append(new_hash)
        return Auth.prover(new_hash, new_tree)
    else:
        new_hash = ctx.consume()
        assert len(new_hash) == 32
        return Auth.verifier(new_hash)


def lookup_auth(
    target: T,
    tree: BST[T] | None,
    ctx: ProofStream,
) -> tuple[T | None, ProofPath]:
    """Authenticated lookup with proof generation.

    Returns (result, proof_path) where proof_path contains
    (direction, sibling_hash, value_hash) along the search path.
    """
    path: ProofPath = []
    node = tree
    while node is not None:
        value_hash = shallow_hash(node.value)
        if target < node.value:  # type: ignore[operator]
            sibling_hash = shallow_hash(node.right) if node.right else LEAF_HASH
            path.append((Direction.LEFT, sibling_hash, value_hash))
            if ctx.mode == Mode.PROVER:
                ctx.append(sibling_hash)
            node = node.left
        elif target > node.value:  # type: ignore[operator]
            sibling_hash = shallow_hash(node.left) if node.left else LEAF_HASH
            path.append((Direction.RIGHT, sibling_hash, value_hash))
            if ctx.mode == Mode.PROVER:
                ctx.append(sibling_hash)
            node = node.right
        else:
            return (node.value, path)
    return (None, path)
