"""Ratchet tree representation and update path processing (RFC 9420 Appendix C).

This module provides a minimal array-indexed tree with helpers for adding,
removing, and updating leaves, computing node hashes, and producing/merging
UpdatePath structures used in commits.
"""

from typing import Any, Optional, Set, cast
from .key_packages import KeyPackage, LeafNode, LeafNodeSource
from ..codec.tls import (
    write_uint8,
    write_uint16,
    write_uint32,
    write_varint,
    write_opaque16,
    read_uint8,
    read_uint16,
    read_opaque16,
    write_opaque_varint,
    read_opaque_varint,
    read_varint,
    read_uint32,
)
from .data_structures import UpdatePath, Signature, serialize_bytes, UpdatePathNode
from . import tree_math
from ..crypto.crypto_provider import CryptoProvider
from ..crypto.hpke_labels import encrypt_with_label, decrypt_with_label
from ..crypto import labels as mls_labels
from ..mls.exceptions import CommitValidationError
import os


class RatchetTreeNode:
    """A single node in the array-indexed ratchet tree.

    Fields
    - is_leaf: True for leaf nodes (even indices), False otherwise.
    - public_key/private_key: Node key material (internal nodes derived from path secrets).
    - parent_hash: Optional parent hash for leaf binding.
    - leaf_node: Serialized leaf node metadata for leaves (KeyPackage.leaf_node).
    - hash: Cached node hash for tree hashing.
    """

    def __init__(self, is_leaf: bool):
        self.is_leaf = is_leaf
        self.public_key: Optional[bytes] = None
        self.private_key: Optional[bytes] = None
        self.parent_hash: Optional[bytes] = None
        self.leaf_node: Optional[LeafNode] = None
        self.hash: Optional[bytes] = None
        # RFC ?7.1: track unmerged leaves for parent nodes
        self.unmerged_leaves: list[int] = [] if not is_leaf else []


class RatchetTree:
    """Array-indexed ratchet tree with hashing and path operations."""

    backend_id = "array"

    def __init__(self, crypto_provider: CryptoProvider):
        """Create an empty ratchet tree."""
        self._n_leaves = 0
        self._nodes: dict[int, RatchetTreeNode] = {}
        self._crypto_provider = crypto_provider

    @property
    def n_leaves(self):
        """Current number of leaves in the tree."""
        return self._n_leaves

    def get_node(self, index: int) -> RatchetTreeNode:
        """Return the node at the given array index, creating a blank node if missing."""
        if index not in self._nodes:
            # Create blank nodes on demand. A node is a leaf if its index is even.
            self._nodes[index] = RatchetTreeNode(index % 2 == 0)
        return self._nodes[index]

    def add_leaf(self, key_package: KeyPackage) -> int:
        """Place new leaf at leftmost blank position (RFC ?12.1.1), or extend tree.

        After placement, all non-blank ancestors on the direct path have the new
        leaf index appended to their unmerged_leaves list (RFC ?7.1).
        """
        if key_package.leaf_node is None:
            raise ValueError("KeyPackage.leaf_node must be present to add a leaf")

        # Find leftmost blank leaf (RFC ?12.1.1)
        leaf_index: Optional[int] = None
        for i in range(self._n_leaves):
            node = self.get_node(i * 2)
            if node.leaf_node is None and node.public_key is None:
                leaf_index = i
                break
        if leaf_index is None:
            # No blank leaf found ? extend tree
            leaf_index = self._n_leaves
            self._n_leaves += 1

        node_index = leaf_index * 2
        node = self.get_node(node_index)
        node.public_key = key_package.leaf_node.encryption_key
        node.leaf_node = key_package.leaf_node

        # RFC ?7.1: add new leaf index to unmerged_leaves of non-blank ancestors
        for ancestor_idx in tree_math.direct_path(node_index, self._n_leaves):
            ancestor = self.get_node(ancestor_idx)
            if ancestor.public_key is not None:
                if leaf_index not in ancestor.unmerged_leaves:
                    ancestor.unmerged_leaves.append(leaf_index)
                    ancestor.unmerged_leaves.sort()

        self._recalculate_hashes_from(node_index)
        return leaf_index

    def remove_leaf(self, index: int) -> None:
        """Blank a leaf and its direct path; then resize to 2^d leaves (RFC ?12.1.3)."""
        node_index = index * 2
        node = self.get_node(node_index)
        node.public_key = None
        node.private_key = None
        node.leaf_node = None

        # Blank the direct path (RFC ?12.1.3)
        for p_idx in tree_math.direct_path(node_index, self.n_leaves):
            p_node = self.get_node(p_idx)
            p_node.public_key = None
            p_node.private_key = None
            p_node.unmerged_leaves = []

        self._recalculate_hashes_from(node_index)

        # RFC ?12.1.3: truncate to smallest 2^d where 2^d >= (rightmost non-blank leaf + 1)
        rightmost = -1
        for i in range(self._n_leaves - 1, -1, -1):
            n = self.get_node(i * 2)
            if n.leaf_node is not None or n.public_key is not None:
                rightmost = i
                break
        if rightmost < 0:
            self._n_leaves = 0
        else:
            # Smallest 2^d strictly greater than rightmost leaf index
            target = 1
            while target <= rightmost:
                target <<= 1
            self._n_leaves = target

    def update_leaf(self, index: int, leaf_node: LeafNode) -> None:
        """Replace leaf metadata and blank direct-path intermediate nodes (RFC ?12.1.2)."""
        node_index = index * 2
        node = self.get_node(node_index)
        node.public_key = leaf_node.encryption_key
        node.leaf_node = leaf_node

        # RFC ?12.1.2: blank all intermediate nodes on the direct path
        for p_idx in tree_math.direct_path(node_index, self.n_leaves):
            p_node = self.get_node(p_idx)
            p_node.public_key = None
            p_node.private_key = None
            p_node.parent_hash = b""
            p_node.unmerged_leaves = []

        self._recalculate_hashes_from(node_index)

    def _compute_parent_hash_for_leaf(self, leaf_index: int) -> bytes:
        """Compute parent hash binding of a leaf to current direct-path nodes (RFC 9420 ?7.9)."""
        if self.n_leaves == 0:
            return b""
        leaf_node_index = leaf_index * 2
        # Ensure hashes are current
        self._recalculate_hashes_from(leaf_node_index)

        path = tree_math.direct_path(leaf_node_index, self.n_leaves)
        parent_hash = b""

        # Iterate from root down to the parent of the leaf
        # Path excludes the leaf itself, so we append it for child lookup.
        reversed_path = list(reversed(path))
        path_nodes_down = reversed_path + [leaf_node_index]

        # We process each node P in the path (which corresponds to reversed_path elements)
        # using the next node in path_nodes_down as the child C.

        for i, p_idx in enumerate(reversed_path):
            node = self.get_node(p_idx)
            child_idx = path_nodes_down[i + 1]

            if not node.public_key:
                # If parent is blank, parent_hash passes through?
                # RFC 9420 ?7.9: "If the node at height h is blank, then the parent hash... is the same as the parent hash of the node at height h+1."
                # effectively, we don't hash it, just keep the accumulator?
                # or do we exclude it from the chain?
                # "The parent hash... is defined as..."
                # If node is blank, it has no `encryption_key`.
                # If blank, we just continue with same parent_hash?
                continue

            # Compute ParentHashInput for this node P
            # We need sibling of C
            sibling_idx = tree_math.sibling(child_idx, self.n_leaves)
            sibling_hash = self._compute_original_sibling_tree_hash(
                parent_index=p_idx,
                sibling_index=sibling_idx,
            )

            # ParentHashInput
            phi = serialize_bytes(node.public_key)
            phi += serialize_bytes(parent_hash)
            phi += serialize_bytes(sibling_hash)

            parent_hash = self._crypto_provider.hash(phi)

        return parent_hash

    def _compute_parent_hash_of_parent_node(
        self, parent_index: int, copath_sibling_index: int
    ) -> bytes:
        """Compute the parent hash of parent node P with copath child S (RFC 9420 ?7.9).

        Returns hash(ParentHashInput) = H(P.encryption_key || P.parent_hash || original_sibling_tree_hash(S)).
        Used to verify that a descendant node D's parent_hash field equals this value.
        """
        node = self.get_node(parent_index)
        if not node.public_key:
            return b""
        sibling_hash = self._compute_original_sibling_tree_hash(parent_index, copath_sibling_index)
        phi = serialize_bytes(node.public_key)
        phi += serialize_bytes(node.parent_hash or b"")
        phi += serialize_bytes(sibling_hash)
        return self._crypto_provider.hash(phi)

    def verify_parent_hash_chains(self) -> None:
        """Verify that every non-blank parent node is parent-hash valid (RFC 9420 ?7.9.2).

        When joining a group, the new member MUST authenticate that each non-blank parent
        node P is parent-hash valid (chainable back to a leaf). Raises CommitValidationError
        if any parent node fails.
        """
        if self.n_leaves == 0:
            return
        root_idx = tree_math.root(self.n_leaves)
        valid_parents: Set[int] = set()
        for leaf_index in range(self.n_leaves):
            leaf_node_index = leaf_index * 2
            node = self.get_node(leaf_node_index)
            if not node.leaf_node or not getattr(node.leaf_node, "parent_hash", None):
                continue
            leaf_ph = node.leaf_node.parent_hash or b""
            if not leaf_ph:
                continue
            try:
                p_idx = tree_math.parent(leaf_node_index, self.n_leaves)
            except Exception:
                continue
            p_node = self.get_node(p_idx)
            if not p_node.public_key:
                continue
            s_idx = tree_math.sibling(leaf_node_index, self.n_leaves)
            expected = self._compute_parent_hash_of_parent_node(p_idx, s_idx)
            if leaf_ph != expected:
                raise CommitValidationError(
                    "parent hash chain invalid: leaf parent_hash does not match computed parent hash"
                )
            # RFC 9420 ?7.9.2 third criterion: D in resolution(C), and
            # intersection(P.unmerged_leaves, subtree(C)) == resolution(C) \ {D}
            d_node = leaf_index * 2
            res_indices = self.resolve_indices(s_idx)
            if d_node not in res_indices:
                raise CommitValidationError(
                    "parent hash invalid: leaf not in resolution of copath sibling"
                )
            res_leaf_indices = {idx // 2 for idx in res_indices if self.get_node(idx).is_leaf}
            leaves_under_s = self._leaf_indices_under(s_idx)
            intersection = set(p_node.unmerged_leaves) & leaves_under_s
            expected_set = res_leaf_indices - {leaf_index}
            if intersection != expected_set:
                raise CommitValidationError(
                    "parent hash invalid: unmerged_leaves intersection with subtree does not match resolution minus D"
                )
            valid_parents.add(p_idx)
            current = p_idx
            while current != root_idx:
                try:
                    p_idx = tree_math.parent(current, self.n_leaves)
                except Exception:
                    break
                p_node = self.get_node(p_idx)
                if not p_node.public_key:
                    break
                s_idx = tree_math.sibling(current, self.n_leaves)
                expected = self._compute_parent_hash_of_parent_node(p_idx, s_idx)
                current_ph = self.get_node(current).parent_hash or b""
                if current_ph != expected:
                    break
                valid_parents.add(p_idx)
                current = p_idx
        for idx in range(self.node_width()):
            n = self.get_node(idx)
            if n.is_leaf:
                continue
            if n.public_key and idx not in valid_parents:
                raise CommitValidationError(
                    f"parent node at index {idx} is not parent-hash valid (RFC 9420 ?7.9.2)"
                )

    def _is_leaf_descendant_of(self, leaf_index: int, ancestor_index: int) -> bool:
        """Return True if the given leaf is contained in ancestor_index's subtree."""
        if self.n_leaves == 0:
            return False
        if leaf_index < 0 or leaf_index >= self.n_leaves:
            return False
        current = leaf_index * 2
        if current == ancestor_index:
            return True
        for p in tree_math.direct_path(current, self.n_leaves):
            if p == ancestor_index:
                return True
        return False

    def _snapshot_node(
        self, node_index: int
    ) -> tuple[
        Optional[bytes],
        Optional[bytes],
        Optional[bytes],
        Optional[LeafNode],
        Optional[bytes],
        list[int],
    ]:
        """Capture node fields so temporary hash modifications can be reverted safely."""
        n = self.get_node(node_index)
        return (
            n.public_key,
            n.private_key,
            n.parent_hash,
            n.leaf_node,
            n.hash,
            list(n.unmerged_leaves),
        )

    def _restore_node(
        self,
        node_index: int,
        snap: tuple[
            Optional[bytes],
            Optional[bytes],
            Optional[bytes],
            Optional[LeafNode],
            Optional[bytes],
            list[int],
        ],
    ) -> None:
        """Restore a node snapshot captured by _snapshot_node."""
        n = self.get_node(node_index)
        n.public_key, n.private_key, n.parent_hash, n.leaf_node, n.hash, n.unmerged_leaves = snap

    def _compute_original_sibling_tree_hash(self, parent_index: int, sibling_index: int) -> bytes:
        """Compute RFC 9420 ?7.9 original_sibling_tree_hash for a parent/copath sibling pair."""
        parent_node = self.get_node(parent_index)
        if not parent_node.unmerged_leaves:
            self._hash_node(sibling_index)
            return self.get_node(sibling_index).hash or b""

        # Collect only unmerged leaves that lie under the sibling subtree.
        leaves_to_blank = [
            leaf_idx
            for leaf_idx in parent_node.unmerged_leaves
            if self._is_leaf_descendant_of(leaf_idx, sibling_index)
        ]
        if not leaves_to_blank:
            self._hash_node(sibling_index)
            return self.get_node(sibling_index).hash or b""

        snapshots: dict[
            int,
            tuple[
                Optional[bytes],
                Optional[bytes],
                Optional[bytes],
                Optional[LeafNode],
                Optional[bytes],
                list[int],
            ],
        ] = {}

        def snapshot_once(idx: int) -> None:
            if idx not in snapshots:
                snapshots[idx] = self._snapshot_node(idx)

        try:
            # Blank each affected leaf and clear references from all parent unmerged sets.
            for leaf_idx in leaves_to_blank:
                leaf_node_index = leaf_idx * 2
                snapshot_once(leaf_node_index)
                leaf_node = self.get_node(leaf_node_index)
                leaf_node.public_key = None
                leaf_node.private_key = None
                leaf_node.leaf_node = None

            for idx, node in list(self._nodes.items()):
                if node.is_leaf or not node.unmerged_leaves:
                    continue
                new_ul = [u for u in node.unmerged_leaves if u not in leaves_to_blank]
                if new_ul != node.unmerged_leaves:
                    snapshot_once(idx)
                    node.unmerged_leaves = new_ul

            # Recompute hashes from each modified leaf to root, then read sibling hash.
            for leaf_idx in leaves_to_blank:
                self._recalculate_hashes_from(leaf_idx * 2)
            self._hash_node(sibling_index)
            return self.get_node(sibling_index).hash or b""
        finally:
            # Restore all temporarily modified nodes.
            for idx, snap in snapshots.items():
                self._restore_node(idx, snap)
            # Refresh hashes along restored paths so caller state remains consistent.
            for leaf_idx in leaves_to_blank:
                self._recalculate_hashes_from(leaf_idx * 2)

    def _recalculate_hashes_from(self, start_node_index: int):
        """Recompute hashes from the given node up to the root along the direct path."""
        self._hash_node(start_node_index)
        path = tree_math.direct_path(start_node_index, self.n_leaves)
        for node_index in path:
            self._hash_node(node_index)

    def _hash_node(self, node_index: int):
        """Compute and cache the hash (NodeHash) for a node using RFC 9420 ?7.8 TreeHashInput."""
        node = self.get_node(node_index)
        if node.is_leaf:
            # TreeHashInput: node_type=1 (leaf)
            blob = write_uint8(1)

            # LeafNodeHashInput structure
            # struct {
            #     uint32 leaf_index;
            #     optional<LeafNode> leaf_node;
            # } LeafNodeHashInput;

            leaf_idx = node_index // 2
            blob += write_uint32(leaf_idx)

            if node.leaf_node:
                blob += write_uint8(1)  # present
                blob += node.leaf_node.serialize()
            else:
                blob += write_uint8(0)  # not present

            node.hash = self._crypto_provider.hash(blob)
        else:
            left_child_hash = self.get_node(tree_math.left(node_index)).hash or b""
            right_child_hash = self.get_node(tree_math.right(node_index, self.n_leaves)).hash or b""

            # TreeHashInput: node_type=2 (parent)
            blob = write_uint8(2)

            # ParentNodeHashInput structure
            # struct {
            #     optional<ParentNode> parent_node;
            #     opaque left_hash<V>;
            #     opaque right_hash<V>;
            # } ParentNodeHashInput;

            if node.public_key:
                # ParentNode present
                blob += write_uint8(1)

                # struct {
                #     HPKEPublicKey encryption_key;
                #     opaque parent_hash<V>;
                #     uint32 unmerged_leaves<V>;
                # } ParentNode;

                blob += serialize_bytes(node.public_key)
                blob += serialize_bytes(node.parent_hash or b"")
                # unmerged_leaves<V>
                ul_data = b""
                for ul in sorted(node.unmerged_leaves):
                    ul_data += write_uint32(ul)
                blob += write_varint(len(ul_data)) + ul_data
            else:
                blob += write_uint8(0)  # not present

            blob += serialize_bytes(left_child_hash)
            blob += serialize_bytes(right_child_hash)

            node.hash = self._crypto_provider.hash(blob)

    def calculate_tree_hash(self) -> bytes:
        """Return the current tree hash (hash of the root), or empty if no leaves."""
        if self.n_leaves == 0:
            return b""
        root_index = tree_math.root(self.n_leaves)
        # Ensure hashes are up to date from root down
        self._hash_node(root_index)
        return self.get_node(root_index).hash or b""

    def resolve(self, node_index: int) -> list[RatchetTreeNode]:
        """
        Return the resolution of a node as defined in RFC 9420 ?4.1.1.
        The resolution is an ordered list of non-blank nodes that cover all non-blank
        descendants of the node.
        """
        node = self.get_node(node_index)
        # If the node is non-blank (has a public key), its resolution is the node
        # itself plus any unmerged leaves.
        if node.public_key is not None:
            res = [node]
            # Append unmerged leaves (RFC ?4.1.1: local unmerged leaves)
            # Note: logic for unmerged leaves in resolution might be complex depending on
            # where they are stored. The RFC says "plus the ordered list of unmerged leaves
            # for the node". Our `unmerged_leaves` field stores indices.
            # We should probably return the nodes for those leaves.
            for leaf_idx in node.unmerged_leaves:
                # Check if the unmerged leaf is actually in the tree
                # (it should be, but let's be safe)
                leaf_node_idx = leaf_idx * 2
                if leaf_node_idx < self.node_width():
                    res.append(self.get_node(leaf_node_idx))
            return res

        # If the node is blank and is a leaf, its resolution is the empty list.
        if node.is_leaf:
            return []

        # If the node is blank and is an intermediate node, its resolution is the
        # concatenation of the resolutions of its children.
        try:
            left_res = self.resolve(tree_math.left(node_index))
            right_res = self.resolve(tree_math.right(node_index, self.n_leaves))
            return left_res + right_res
        except Exception:
            # If children don't exist (e.g. right child out of bounds), treat as empty
            return []

    def resolve_indices(self, node_index: int) -> list[int]:
        """Return the resolution of a node as a list of node indices (RFC 9420 ?4.1.1)."""
        node = self.get_node(node_index)
        if node.public_key is not None:
            res = [node_index]
            for leaf_idx in node.unmerged_leaves:
                leaf_node_idx = leaf_idx * 2
                if leaf_node_idx < self.node_width():
                    res.append(leaf_node_idx)
            return res
        if node.is_leaf:
            return []
        try:
            left_res = self.resolve_indices(tree_math.left(node_index))
            right_res = self.resolve_indices(tree_math.right(node_index, self.n_leaves))
            return left_res + right_res
        except Exception:
            return []

    def _leaf_indices_under(self, node_index: int) -> Set[int]:
        """Return set of leaf indices that are descendants of the given node."""
        out: Set[int] = set()
        node = self.get_node(node_index)
        if node.is_leaf:
            out.add(node_index // 2)
            return out
        try:
            left_idx = tree_math.left(node_index)
            right_idx = tree_math.right(node_index, self.n_leaves)
            out |= self._leaf_indices_under(left_idx)
            out |= self._leaf_indices_under(right_idx)
        except Exception:
            pass
        return out

    def filtered_direct_path(self, node_index: int) -> list[int]:
        """
        Return the filtered direct path of a node (RFC 9420 ?4.1.2).
        The filtered direct path is the direct path with nodes removed if their
        child on the copath has an empty resolution.
        """
        d = tree_math.direct_path(node_index, self.n_leaves)
        filtered: list[int] = []
        for p in d:
            # Child on the direct path is the one from which we came.
            # We need to find the child of p that is on the copath of node_index.
            # Actually, the definition says: "remove nodes whose child on the *copath*
            # has empty resolution".
            # The children of p are (left, right). One is on the direct path, the other is on the copath.
            # We check the resolution of the *sibling* of the node on the direct path.

            # Find the child of p that is on the direct path (or is node_index itself)
            # We can iterate up.
            # But simpler: the child on the copath is simply the sibling of the child on the
            # direct path.
            # Wait, `p` is in the direct path of `node_index`.
            # Let `c` be the child of `p` that is also in the direct path (or `node_index`).
            # The "child on the copath" is `sibling(c)`.

            # Let's find `c`. `c` is the node below `p` in the path.
            # Since `d` is ordered from leaf to root (in our tree_math, let's verify),
            # `direct_path` returns "indices on the path from node x up to ... root".
            # So `d[0]` is parent of `node_index`. `d[1]` is parent of `d[0]`.

            if not filtered and p == tree_math.parent(node_index, self.n_leaves):
                c = node_index
            else:
                # If we have started building filtered, p is parent of the last added node?
                # No, d is independent.
                # We need to look at the previous node in the chain x -> ... -> root.
                # If d = [p1, p2, ...], x -> p1 -> p2.
                # For p1, child is x.
                # For p2, child is p1.
                if d.index(p) == 0:
                    c = node_index
                else:
                    c = d[d.index(p) - 1]

            s = tree_math.sibling(c, self.n_leaves)

            # Check resolution of s
            if self.resolve(s):
                filtered.append(p)

        return filtered

    def node_width(self) -> int:
        return tree_math.node_width(self.n_leaves)

    def create_update_path(
        self,
        committer_index: int,
        new_leaf_node: LeafNode,
        group_context_bytes: bytes,
        excluded_leaf_pubkeys: Optional[Set[bytes]] = None,
    ) -> tuple[UpdatePath, bytes, dict[int, bytes]]:
        """Create an UpdatePath for the committer and derive the commit secret.

        RFC 9420 ?7.4: Generate a fresh path_secret at the leaf and derive
        subsequent path/node secrets top-down up the direct path. For each
        direct-path node, derive a deterministic key pair from the node_secret.
        Encrypt the path_secret for the copath resolution. Attach a parent hash
        to the new leaf node. Return (UpdatePath, commit_secret), where
        commit_secret is the final path_secret at the root.

        Args:
            excluded_leaf_pubkeys: HPKE public keys of newly added members that
                MUST be excluded from path-secret encryption (RFC ?7.4). These
                members were added in the same commit and cannot decrypt path
                secrets because they don't yet have the tree secrets.
        """
        direct_path = self.filtered_direct_path(committer_index * 2)
        full_direct_path = tree_math.direct_path(committer_index * 2, self.n_leaves)
        filtered_set = set(direct_path)

        if not full_direct_path:
            # Single-member tree
            update_path = UpdatePath(new_leaf_node.serialize(), [])
            commit_secret = self._crypto_provider.kdf_extract(b"", b"")
            return update_path, commit_secret, {}

        # RFC ?7.4: path_secret[0] = random; use current, THEN derive next (use-then-derive)
        current_path_secret = os.urandom(self._crypto_provider.kdf_hash_len())

        # For each node on the direct path, use current secret for this node, then advance
        path_secret_by_node: dict[int, bytes] = {}

        for node_index in full_direct_path:
            # Use current_path_secret for this node BEFORE advancing
            if node_index in filtered_set:
                path_secret_by_node[node_index] = current_path_secret
                node_secret = self._crypto_provider.derive_secret(current_path_secret, b"node")
                priv_key, pub_key = self._crypto_provider.derive_key_pair(node_secret)
                self._nodes[node_index].private_key = priv_key
                self._nodes[node_index].public_key = pub_key
                self._nodes[node_index].unmerged_leaves = []  # cleared by committer
            else:
                self._nodes[node_index].private_key = None
                self._nodes[node_index].public_key = None
                self._nodes[node_index].unmerged_leaves = []
                self._nodes[node_index].leaf_node = None
            # Advance to next path_secret for the next node
            current_path_secret = self._crypto_provider.derive_secret(current_path_secret, b"path")

        update_path_nodes: list[UpdatePathNode] = []

        for node_index in direct_path:
            node = self.get_node(node_index)
            if not node.public_key:
                # Should not happen if filtered correctly
                continue
                # raise ValueError(f"Node {node_index} in filtered path but has no public key")

            # Encrypt path secret to resolution of copath node
            try:
                copath_node_index = tree_math.sibling(node_index, self.n_leaves)
            except Exception:
                copath_node_index = None

            encrypted_secrets = []
            if copath_node_index is not None:
                path_secret = path_secret_by_node[node_index]
                resolution_nodes = self.resolve(copath_node_index)

                # RFC ?7.4: exclude newly-added leaves ? they were added in this
                # commit and have no tree secrets to decrypt path secrets.
                if excluded_leaf_pubkeys:
                    resolution_nodes = [
                        rn for rn in resolution_nodes if rn.public_key not in excluded_leaf_pubkeys
                    ]

                for res_node in resolution_nodes:
                    if res_node.public_key:
                        enc, ct = encrypt_with_label(
                            self._crypto_provider,
                            recipient_public_key=res_node.public_key,
                            label=mls_labels.HPKE_UPDATE_PATH_NODE,
                            context=group_context_bytes,
                            aad=b"",
                            plaintext=path_secret,
                        )
                        hpke_ciphertext = write_opaque_varint(enc) + write_opaque_varint(ct)
                        encrypted_secrets.append(hpke_ciphertext)

            update_path_nodes.append(
                UpdatePathNode(
                    encryption_key=node.public_key, encrypted_path_secrets=encrypted_secrets
                )
            )

        parent_hash = self._compute_parent_hash_for_leaf(committer_index)

        # Create leaf with COMMIT source
        leaf_for_path = LeafNode(
            encryption_key=new_leaf_node.encryption_key,
            signature_key=new_leaf_node.signature_key,
            credential=new_leaf_node.credential,
            capabilities=new_leaf_node.capabilities,
            leaf_node_source=LeafNodeSource.COMMIT,
            parent_hash=parent_hash,
            lifetime_not_before=new_leaf_node.lifetime_not_before,
            lifetime_not_after=new_leaf_node.lifetime_not_after,
            extensions=new_leaf_node.extensions,
            signature=new_leaf_node.signature,
        )

        update_path = UpdatePath(leaf_for_path.serialize(), update_path_nodes)
        # RFC ?7.4: commit_secret = path_secret[root], i.e. the path_secret for the
        # root node, NOT the one advanced past the root after the loop ends.
        root_node_index = full_direct_path[-1]
        commit_secret = path_secret_by_node.get(root_node_index, current_path_secret)
        # RFC 9420 ?12.4.3: Return path_secret_by_node so Welcome can include path_secret for each joiner's LCA.
        return update_path, commit_secret, path_secret_by_node

    def merge_update_path(
        self, update_path: UpdatePath, committer_index: int, group_context_bytes: bytes
    ) -> bytes:
        """Merge an UpdatePath from a received commit and return the commit secret.

        RFC 9420 ?7.4 receive path:
        - Verify parent hash binding for provided leaf (if present)
        - Decrypt exactly one path_secret corresponding to a copath node on our
          direct path; then derive subsequent path/node secrets upward
        - Update keys along the direct path and recompute hashes
        """
        # Update leaf node
        provided_leaf = LeafNode.deserialize(update_path.leaf_node)
        self.update_leaf(committer_index, provided_leaf)

        # Decrypt a single path_secret for the lowest applicable node; then derive upwards
        # The UpdatePath nodes correspond to the *filtered direct path*.
        # We need to map them to our direct path.

        # RFC 9420 ?7.4:
        # "The receiver... identifies the first node in the filtered direct path... for which it possesses a private key in the resolution of the copath node."

        filtered_path = self.filtered_direct_path(committer_index * 2)
        # RFC 9420 ?12.4.2: path encryption keys MUST be unique in the tree.
        try:
            from .validations import validate_update_path_key_uniqueness

            validate_update_path_key_uniqueness(
                self,
                [n.encryption_key for n in update_path.nodes if n.encryption_key],
                committer_index,
            )
        except Exception as e:
            if isinstance(e, CommitValidationError):
                raise
            raise CommitValidationError(f"invalid UpdatePath encryption keys: {e}") from e

        decrypted_index: Optional[int] = None
        current_path_secret: Optional[bytes] = None

        # Iterate over both the filtered path and the UpdatePath nodes
        # They should align.
        if len(update_path.nodes) != len(filtered_path):
            # This might happen if our view of filtered path differs?
            # Or if the sender's filtered path excluded something correctly that we included, or vice versa.
            # RFC says "The UpdatePath... contains a list of UpdatePathNode... corresponding to the filtered direct path".
            # If we disagree on filtered path, we might misinterpret.
            # However, we can try to find a match.
            pass

        for i, node_index in enumerate(filtered_path):
            if i >= len(update_path.nodes):
                break

            up_node = update_path.nodes[i]

            # Update the node's public key; clear unmerged_leaves (RFC ?7.5)
            nd = self.get_node(node_index)
            nd.public_key = up_node.encryption_key
            nd.unmerged_leaves = []

            if current_path_secret is not None:
                # We already have the secret, just updating keys
                continue

            # Try to decrypt
            # Skip root (no encrypted secrets)
            if node_index == tree_math.root(self.n_leaves):
                continue

            # Check if we can decrypt
            # We need to be in the resolution of the copath node.
            # Do we have a private key in the resolution of the copath node?
            # Wait, we just need to try decrypting the blobs.
            # There is a list of blobs. One of them might be for us.

            # We are the receiver. We have our own leaf key, and potentially other node keys.
            # We should try to decrypt with any private key we have that is in the resolution of the copath node?
            # Actually, the sender encrypted to the resolution.
            # We just need to check if any of our private keys can decrypt any of the cyphertexts.

            # Which private keys do we possess?
            # We possess keys for nodes on our direct path.
            # The copath node is a sibling of a node on the committer's direct path.
            # If we are in the subtree of the copath node, we have keys.

            # Optimization:
            # 1. Identify valid private keys we hold.
            # 2. Try to decrypt.

            # In this simple implementation, we can iterate our known private keys?
            # Or simpler:
            # The sender encrypted to specific public keys.
            # Those public keys belong to nodes in the tree.
            # If we hold the private key for a node, we can try.

            # But the ciphertext doesn't say "this is for node X". It's just a list.
            # We have to try all our keys against all ciphertexts? That's O(M*N).
            # RFC says "The position in the list... corresponds to the position in the resolution".
            # So we need to compute the resolution of the copath node (same as sender did).
            # Then we check if we hold the private key for the node at index J in the resolution.
            # If so, we attempt to decrypt the J-th ciphertext.

            copath_node_index = tree_math.sibling(node_index, self.n_leaves)
            resolution = self.resolve(copath_node_index)
            # RFC 9420 ?7.6: Reject when encrypted_path_secret count does not match resolution count.
            if len(up_node.encrypted_path_secrets) != len(resolution):
                raise CommitValidationError(
                    "UpdatePath node encrypted_path_secret count does not match copath resolution length (RFC 9420 ?7.6)"
                )

            for j, res_node in enumerate(resolution):
                # Do we have the private key for `res_node`?
                # `res_node` is a RatchetTreeNode object from `self.get_node`.
                if res_node.private_key:
                    # Try decrypting the j-th secret
                    blob = up_node.encrypted_path_secrets[j]
                    try:
                        from .data_structures import deserialize_bytes
                        # blob is HPKECiphertext (kem_output || ciphertext), but serialized as bytes?
                        # In UpdatePathNode.deserialize we stored raw bytes.
                        # HPKECiphertext = opaque kem<V>; opaque ct<V>;
                        # We need to parse it.

                        kem_out, rest_ct = deserialize_bytes(blob)
                        ct, _ = deserialize_bytes(rest_ct)

                        ps = decrypt_with_label(
                            self._crypto_provider,
                            recipient_private_key=res_node.private_key,
                            kem_output=kem_out,
                            label=mls_labels.HPKE_UPDATE_PATH_NODE,
                            context=group_context_bytes,
                            aad=b"",
                            ciphertext=ct,
                        )
                        decrypted_index = node_index
                        current_path_secret = ps
                        break
                    except Exception:
                        continue

            if current_path_secret is not None:
                # Found it!
                pass  # Continue loop to update remaining public keys

        if current_path_secret is None:
            # RFC 9420 ?7.6: decryption failure MUST cause the commit to be rejected
            raise CommitValidationError(
                "merge_update_path: could not decrypt any path secret from UpdatePath"
            )

        # From the decrypted node onward, derive subsequent path/node secrets upward
        # We need to find where `decrypted_index` is in the full direct path.
        full_direct_path = tree_math.direct_path(committer_index * 2, self.n_leaves)

        start_idx = (
            full_direct_path.index(decrypted_index) if decrypted_index in full_direct_path else 0
        )

        for node_index in full_direct_path[start_idx:]:
            # For the first node, use current_path_secret as decrypted; otherwise, step the ratchet
            if node_index != decrypted_index:
                current_path_secret = self._crypto_provider.derive_secret(
                    current_path_secret, b"path"
                )

            # Derive keys for this node
            node_secret = self._crypto_provider.derive_secret(current_path_secret, b"node")
            priv_key, pub_key = self._crypto_provider.derive_key_pair(node_secret)

            # RFC ?7.4: MUST verify that the derived key matches the key in UpdatePath.
            # The UpdatePath node for this index (if it's in filtered_path)
            if node_index in filtered_path:
                up_idx = filtered_path.index(node_index)
                if up_idx < len(update_path.nodes):
                    expected_pk = update_path.nodes[up_idx].encryption_key
                    if expected_pk and pub_key != expected_pk:
                        raise CommitValidationError(
                            f"Derived public key at node {node_index} does not match UpdatePath"
                        )

            # Update the node ? we have the private key for decrypted/derived nodes.
            self.get_node(node_index).private_key = priv_key
            self.get_node(node_index).public_key = pub_key

        self._recalculate_hashes_from(committer_index * 2)
        # RFC 9420 ?7.5: Compute and store parent_hash on each node of the sender's filtered direct path.
        for path_node_index in filtered_path:
            try:
                parent_idx = tree_math.parent(path_node_index, self.n_leaves)
                s_idx = tree_math.sibling(path_node_index, self.n_leaves)
                ph = self._compute_parent_hash_of_parent_node(parent_idx, s_idx)
                self.get_node(parent_idx).parent_hash = ph if ph else None
            except Exception:
                continue
        # Re-verify parent hash after applying path secrets to ensure consistency
        try:
            if provided_leaf.parent_hash:
                expected_after = self._compute_parent_hash_for_leaf(committer_index)
                if expected_after != provided_leaf.parent_hash:
                    # Warn or fail? RFC says MUST verify.
                    raise CommitValidationError("parent_hash mismatch after applying update path")
        except Exception:
            # If hash computation fails or mismatch
            raise CommitValidationError("parent_hash validation failed")

        # Commit secret is the final path_secret at the root of the direct path
        commit_secret = current_path_secret
        return commit_secret

    def apply_joiner_path_secret(self, leaf_index: int, path_secret_at_leaf: bytes) -> None:
        """Apply path_secret from GroupSecrets when joining via Welcome (RFC 9420 ?12.4.3.1).

        Derives path secrets from the joiner's leaf up to the root and stores the
        resulting private keys in the tree. Call after loading the tree from the
        ratchet_tree extension and before deriving epoch secrets.
        """
        full_direct_path = tree_math.direct_path(leaf_index * 2, self.n_leaves)
        current = path_secret_at_leaf
        for node_index in full_direct_path:
            node_secret = self._crypto_provider.derive_secret(current, b"node")
            priv_key, pub_key = self._crypto_provider.derive_key_pair(node_secret)
            self.get_node(node_index).private_key = priv_key
            self.get_node(node_index).public_key = pub_key
            current = self._crypto_provider.derive_secret(current, b"path")
        self._recalculate_hashes_from(leaf_index * 2)

    # --- Welcome ratchet_tree extension helpers ---
    def serialize_tree_for_welcome(self) -> bytes:
        """
        Serialize leaves needed for a new member to reconstruct the tree view.
        Format:
          uint16 n_leaves
          repeated opaque16 leaf_node (empty for blank leaves)
        """
        out = write_uint16(self.n_leaves)
        for leaf in range(self.n_leaves):
            node = self.get_node(leaf * 2)
            if node.leaf_node:
                out += write_opaque16(node.leaf_node.serialize())
            else:
                out += write_opaque16(b"")
        return out

    def serialize_full_tree_for_welcome(self) -> bytes:
        """
        Serialize the ratchet tree per RFC 9420 ?12.4.3.1: optional<Node> ratchet_tree<V>.

        Nodes are in array order (left-to-right in-order). Blank nodes after the
        last non-blank node are omitted per RFC. Each optional<Node> is encoded
        as varint(0) for blank or varint(1) + Node (NodeType + LeafNode or ParentNode).
        """
        width = tree_math.node_width(self.n_leaves)
        # Find last non-blank index (RFC: sender MUST NOT include blank after last non-blank)
        last_non_blank = -1
        for idx in range(width):
            node = self.get_node(idx)
            if node.is_leaf:
                if node.leaf_node is not None:
                    last_non_blank = idx
            else:
                if node.public_key or node.parent_hash:
                    last_non_blank = idx
        if last_non_blank < 0:
            return write_varint(0)  # empty vector
        content = b""
        for idx in range(last_non_blank + 1):
            node = self.get_node(idx)
            if node.is_leaf:
                if node.leaf_node:
                    # optional present: varint(1) + Node(leaf(1) + LeafNode)
                    content += write_varint(1)
                    content += write_uint8(1)  # NodeType.leaf
                    content += write_opaque_varint(node.leaf_node.serialize())
                else:
                    content += write_varint(0)
            else:
                if node.public_key or node.parent_hash:
                    content += write_varint(1)
                    content += write_uint8(2)  # NodeType.parent
                    # ParentNode: encryption_key, parent_hash, unmerged_leaves<V>
                    content += serialize_bytes(node.public_key or b"")
                    content += serialize_bytes(node.parent_hash or b"")
                    ul_data = b"".join(write_uint32(ul) for ul in sorted(node.unmerged_leaves))
                    content += write_varint(len(ul_data)) + ul_data
                else:
                    content += write_varint(0)
        return write_varint(len(content)) + content

    def load_tree_from_welcome_bytes(self, data: bytes) -> None:
        """
        Load leaves from a ratchet_tree extension as serialized by serialize_tree_for_welcome().
        """
        off = 0
        n, off = read_uint16(data, off)
        self._n_leaves = 0
        self._nodes.clear()
        for i in range(n):
            blob, off = read_opaque16(data, off)
            if blob:
                leaf = LeafNode.deserialize(blob)
                self.add_leaf(
                    KeyPackage(leaf_node=leaf, signature=Signature(b""))
                )  # signature not validated here
            else:
                # Even if blank, we need to advance the leaf count
                self._n_leaves += 1
        # Re-hash tree
        if self.n_leaves > 0:
            self._recalculate_hashes_from(0)

    def load_full_tree_from_welcome_bytes(self, data: bytes) -> None:
        """
        Load the tree from RFC 9420 ?12.4.3.1 optional<Node> ratchet_tree<V>.

        The receiver extends with blank nodes until length is 2^(d+1)-1, then
        infers n_leaves = 2^d.
        """
        from ..codec.tls import TLSDecodeError

        off = 0
        total_len, off = read_varint(data, off)
        if total_len == 0:
            self._n_leaves = 0
            self._nodes.clear()
            return
        end = off + total_len
        if end > len(data):
            raise TLSDecodeError("ratchet_tree extension shorter than declared length")
        # (node_type, payload): payload is bytes for leaf (1), (pk, ph, unmerged_leaves) for parent (2)
        entries: list[Optional[tuple[int, Any]]] = []
        while off < end:
            opt_len, off = read_varint(data, off)
            if opt_len == 0:
                entries.append(None)
                continue
            if opt_len != 1:
                raise TLSDecodeError("optional<Node> must have length 0 or 1")
            node_type, off = read_uint8(data, off)
            if node_type == 1:  # leaf
                leaf_bytes, off = read_opaque_varint(data, off)
                entries.append((1, leaf_bytes))
            elif node_type == 2:  # parent: encryption_key, parent_hash, unmerged_leaves<V>
                pk, off = read_opaque_varint(data, off)
                ph, off = read_opaque_varint(data, off)
                ul_len, off = read_varint(data, off)
                unmerged_leaves: list[int] = []
                ul_end = off + ul_len
                while off < ul_end and off + 4 <= len(data):
                    ul_val, off = read_uint32(data, off)
                    unmerged_leaves.append(ul_val)
                entries.append((2, (pk, ph, unmerged_leaves)))
            else:
                raise TLSDecodeError(f"invalid NodeType {node_type}")
        # RFC: extend with blank until length = 2^(d+1)-1
        n_entries = len(entries)
        if n_entries == 0:
            self._n_leaves = 0
            self._nodes.clear()
            return
        # Minimum width = 2^(d+1)-1 >= n_entries => d = ceil(log2(n_entries+1)) - 1
        d = 0
        while (1 << (d + 1)) - 1 < n_entries:
            d += 1
        width = (1 << (d + 1)) - 1
        n_leaves = 1 << d
        self._n_leaves = n_leaves
        self._nodes.clear()
        for i in range(width):
            self._nodes[i] = RatchetTreeNode(is_leaf=(i % 2 == 0))
        # Fill from entries (entries may be shorter; rest are blank)
        off = 0
        for idx in range(width):
            if idx >= len(entries):
                continue
            ent = entries[idx]
            if ent is None:
                continue
            node_type, payload = ent
            node = self.get_node(idx)
            if node_type == 1:
                leaf = LeafNode.deserialize(cast(bytes, payload))
                node.public_key = leaf.encryption_key
                node.leaf_node = leaf
            else:
                pk, ph, unmerged_leaves = cast(tuple[Any, Any, Any], payload)
                node.public_key = pk if pk else None
                node.parent_hash = ph if ph else None
                node.unmerged_leaves = list(unmerged_leaves) if unmerged_leaves else []
        if self.n_leaves > 0:
            self._recalculate_hashes_from(0)

    # --- Persistence helpers (full state, including private keys when present) ---
    def serialize_full_state(self) -> bytes:
        """
        Serialize the full ratchet tree state, including private keys where present.
        Format:
          uint16 n_leaves
          uint16 node_count
          repeated {
            uint8 node_type (0=blank, 1=leaf, 2=parent)
            if leaf:
              opaque16 leaf_node (may be empty)
              opaque16 public_key (may be empty)
              opaque16 private_key (may be empty)
            if parent:
              opaque16 public_key (may be empty)
              opaque16 private_key (may be empty)
              opaque16 parent_hash (may be empty)
          }
        """
        out = write_uint16(self.n_leaves)
        width = tree_math.node_width(self.n_leaves)
        out += write_uint16(width)
        for idx in range(width):
            node = self.get_node(idx)
            if node.is_leaf:
                if node.leaf_node or node.public_key or node.private_key:
                    out += write_uint8(1)
                    out += write_opaque16(node.leaf_node.serialize() if node.leaf_node else b"")
                    out += write_opaque16(node.public_key or b"")
                    out += write_opaque16(node.private_key or b"")
                else:
                    out += write_uint8(0)
            else:
                if node.public_key or node.private_key or node.parent_hash:
                    out += write_uint8(2)
                    out += write_opaque16(node.public_key or b"")
                    out += write_opaque16(node.private_key or b"")
                    out += write_opaque16(node.parent_hash or b"")
                else:
                    out += write_uint8(0)
        return out

    def load_full_state(self, data: bytes) -> None:
        """
        Load a full ratchet tree state produced by serialize_full_state().
        """
        off = 0
        n, off = read_uint16(data, off)
        width, off = read_uint16(data, off)
        self._n_leaves = 0
        self._nodes.clear()
        for _ in range(n):
            self._n_leaves += 1
        for idx in range(width):
            node_type, off = read_uint8(data, off)
            if node_type == 0:
                continue
            node = self.get_node(idx)
            if node.is_leaf:
                blob, off = read_opaque16(data, off)
                pk, off = read_opaque16(data, off)
                sk, off = read_opaque16(data, off)
                if blob:
                    leaf = LeafNode.deserialize(blob)
                    node.leaf_node = leaf
                    node.public_key = leaf.encryption_key
                else:
                    node.public_key = pk if pk else None
                node.private_key = sk if sk else None
            else:
                pk, off = read_opaque16(data, off)
                sk, off = read_opaque16(data, off)
                ph, off = read_opaque16(data, off)
                node.public_key = pk if pk else None
                node.private_key = sk if sk else None
                node.parent_hash = ph if ph else None
        if self.n_leaves > 0:
            self._recalculate_hashes_from(0)
