"""Merkle tree helpers for AMCS memory roots."""

from __future__ import annotations

from amcs.hashing import sha256_hex

EMPTY_MERKLE_ROOT_HEX = sha256_hex(b"")


def _pair_hash(left_hex: str, right_hex: str) -> str:
    return sha256_hex(bytes.fromhex(left_hex) + bytes.fromhex(right_hex))


def merkle_root_hex(leaf_hashes_hex: list[str]) -> str:
    """Compute a Merkle root from ordered leaf hashes (hex SHA-256 strings)."""
    if not leaf_hashes_hex:
        return EMPTY_MERKLE_ROOT_HEX

    level = list(leaf_hashes_hex)
    while len(level) > 1:
        if len(level) % 2 == 1:
            level.append(level[-1])

        next_level: list[str] = []
        for i in range(0, len(level), 2):
            next_level.append(_pair_hash(level[i], level[i + 1]))
        level = next_level

    return level[0]


def build_inclusion_proof(leaf_hashes: list[str], index: int) -> list[tuple[str, str]]:
    """Build a Merkle inclusion proof for the leaf at index."""
    if not leaf_hashes:
        raise ValueError("leaf_hashes must not be empty")
    if index < 0 or index >= len(leaf_hashes):
        raise IndexError("index out of range")

    level = list(leaf_hashes)
    idx = index
    proof: list[tuple[str, str]] = []

    while len(level) > 1:
        if len(level) % 2 == 1:
            level.append(level[-1])

        sibling_idx = idx - 1 if idx % 2 == 1 else idx + 1
        direction = "left" if sibling_idx < idx else "right"
        proof.append((direction, level[sibling_idx]))

        next_level: list[str] = []
        for i in range(0, len(level), 2):
            next_level.append(_pair_hash(level[i], level[i + 1]))

        level = next_level
        idx //= 2

    return proof


def verify_inclusion_proof(leaf_hash: str, proof: list[tuple[str, str]], root: str) -> bool:
    """Verify a Merkle inclusion proof for a leaf hash and expected root."""
    current = leaf_hash
    for direction, sibling in proof:
        if direction == "left":
            current = _pair_hash(sibling, current)
        elif direction == "right":
            current = _pair_hash(current, sibling)
        else:
            return False

    return current == root
