"""Balanced binary Merkle tree with membership proofs.

Ports ADS4All's Merkle/Binary.hs: build from list, index-based proof
generation, verified reconstruction, and authenticated insert/lookup.

Domain separation (RFC 9162 pattern):
  - Leaf:     combine_hashes(sha256(b"MerkleLeaf"), shallow_hash(value))
  - Internal: combine_hashes(sha256(b"MerkleNode"), left.hash, right.hash)
  - Empty:    sha256(b"EmptyTree")
"""

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")

EMPTY_TREE_HASH: bytes = sha256(b"EmptyTree")
"""Sentinel hash for empty Merkle trees."""

_LEAF_TAG: bytes = sha256(b"MerkleLeaf")
_NODE_TAG: bytes = sha256(b"MerkleNode")


class Side(Enum):
    """Which side the sibling is on in a proof path entry."""

    LEFT = auto()
    RIGHT = auto()


@dataclass(frozen=True, slots=True)
class MerkleNode:
    """Binary Merkle tree node. Leaves carry data; internal nodes carry children."""

    hash: bytes
    left: MerkleNode | None = None
    right: MerkleNode | None = None
    data: object | None = None  # Only at leaves

    def __ads_hash__(self) -> bytes:
        return self.hash


@dataclass(frozen=True, slots=True)
class MerkleProof:
    """Membership proof: path of (side, sibling_hash) pairs plus root hash."""

    path: list[tuple[Side, bytes]]
    root_hash: bytes


# ---------------------------------------------------------------------------
# Building
# ---------------------------------------------------------------------------


def _make_leaf(value: object) -> MerkleNode:
    """Create a leaf node with domain-separated hash."""
    h = combine_hashes(_LEAF_TAG, shallow_hash(value))
    return MerkleNode(hash=h, data=value)


def _make_internal(left: MerkleNode, right: MerkleNode) -> MerkleNode:
    """Create an internal node with domain-separated hash."""
    h = combine_hashes(_NODE_TAG, left.hash, right.hash)
    return MerkleNode(hash=h, left=left, right=right)


def build_merkle(data: list) -> MerkleNode | None:
    """Build a balanced binary Merkle tree from a list.

    Empty list -> None.  Single element -> leaf.
    Otherwise split at midpoint and recurse.
    """
    if not data:
        return None
    if len(data) == 1:
        return _make_leaf(data[0])
    mid = len(data) // 2
    left = build_merkle(data[:mid])
    right = build_merkle(data[mid:])
    # Both are non-None because each half has at least one element
    assert left is not None and right is not None
    return _make_internal(left, right)


# ---------------------------------------------------------------------------
# Queries
# ---------------------------------------------------------------------------


def root_hash(tree: MerkleNode | None) -> bytes:
    """Return the root hash, or EMPTY_TREE_HASH for None."""
    if tree is None:
        return EMPTY_TREE_HASH
    return tree.hash


def size(tree: MerkleNode | None) -> int:
    """Count the number of leaves."""
    if tree is None:
        return 0
    if tree.data is not None:
        return 1
    return size(tree.left) + size(tree.right)


def height(tree: MerkleNode | None) -> int:
    """Tree depth (0 for None, 1 for leaf)."""
    if tree is None:
        return 0
    if tree.data is not None:
        return 1
    return 1 + max(height(tree.left), height(tree.right))


def leaves(tree: MerkleNode | None) -> list:
    """Return leaf data in left-to-right order."""
    if tree is None:
        return []
    if tree.data is not None:
        return [tree.data]
    return leaves(tree.left) + leaves(tree.right)


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


def generate_proof(index: int, tree: MerkleNode, data_size: int) -> MerkleProof:
    """Generate a membership proof for the leaf at `index`.

    `data_size` is the number of leaves in the tree (needed to compute
    the split point at each level).

    Returns a MerkleProof with a path of (Side, sibling_hash) entries
    from root to leaf.
    """
    if data_size <= 0:
        raise ValueError("Cannot generate proof for empty tree")
    if index < 0 or index >= data_size:
        raise IndexError(f"Index {index} out of range [0, {data_size})")

    path: list[tuple[Side, bytes]] = []
    node = tree
    lo, hi = 0, data_size

    while hi - lo > 1:
        mid = lo + (hi - lo) // 2
        assert node.left is not None and node.right is not None
        if index < mid:
            # Go left; sibling is on the right
            path.append((Side.RIGHT, node.right.hash))
            node = node.left
            hi = mid
        else:
            # Go right; sibling is on the left
            path.append((Side.LEFT, node.left.hash))
            node = node.right
            lo = mid

    return MerkleProof(path=path, root_hash=tree.hash)


def verify_proof(value: object, proof: MerkleProof) -> bool:
    """Verify a membership proof by reconstructing root from leaf upward.

    Uses hmac.compare_digest for constant-time comparison.
    """
    current = combine_hashes(_LEAF_TAG, shallow_hash(value))

    for side, sibling_hash in reversed(proof.path):
        if side == Side.LEFT:
            # Sibling is on the left, we are on the right
            current = combine_hashes(_NODE_TAG, sibling_hash, current)
        else:
            # Sibling is on the right, we are on the left
            current = combine_hashes(_NODE_TAG, current, sibling_hash)

    return hmac.compare_digest(current, proof.root_hash)


# ---------------------------------------------------------------------------
# Authenticated operations (follow BST pattern)
# ---------------------------------------------------------------------------


def insert_auth(
    x: object,
    auth_tree: Auth,
    ctx: ProofStream,
) -> Auth:
    """Authenticated insert: rebuild tree with new element appended.

    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()
        old_leaves = leaves(tree)
        new_tree = build_merkle(old_leaves + [x])
        new_hash = root_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 lookup_auth(
    index: int,
    tree: MerkleNode,
    data_size: int,
    ctx: ProofStream,
) -> MerkleProof:
    """Authenticated lookup: generate proof and emit sibling hashes.

    Prover: generates proof, appends sibling hashes to proof stream.
    Verifier: consumes sibling hashes from proof stream.
    Returns the MerkleProof for standalone verification.
    """
    proof = generate_proof(index, tree, data_size)
    if ctx.mode == Mode.PROVER:
        for _side, sibling_hash in proof.path:
            ctx.append(sibling_hash)
    else:
        for _side, _sibling_hash in proof.path:
            ctx.consume()
    return proof
