"""zg_storage.transfer.merkle module."""

from __future__ import annotations

from dataclasses import dataclass
from typing import List

from web3 import Web3


@dataclass
class Node:
    """Node class.

    Purpose:
        SDK component type."""

    hash: bytes
    parent: "Node | None" = None
    left: "Node | None" = None
    right: "Node | None" = None

    def is_left(self) -> bool:
        return self.parent is not None and self.parent.left is self


class MerkleTree:
    """MerkleTree class.

    Purpose:
        SDK component type."""

    def __init__(self, leaves: List[bytes]) -> None:
        if not leaves:
            raise ValueError("no leaves")
        self.leaf_nodes = [Node(h) for h in leaves]
        self.root = self._build(self.leaf_nodes)

    def _build(self, leaves: List[Node]) -> Node:
        """Build the tree layers upward from leaf nodes."""
        queue: List[Node] = []
        i = 0
        while i < len(leaves):
            if i == len(leaves) - 1:
                queue.append(leaves[i])
                i += 1
                continue
            left = leaves[i]
            right = leaves[i + 1]
            parent = Node(hash=Web3.keccak(left.hash + right.hash), left=left, right=right)
            left.parent = parent
            right.parent = parent
            queue.append(parent)
            i += 2

        while len(queue) > 1:
            next_level: List[Node] = []
            i = 0
            while i < len(queue):
                if i == len(queue) - 1:
                    next_level.append(queue[i])
                    i += 1
                    continue
                left = queue[i]
                right = queue[i + 1]
                parent = Node(hash=Web3.keccak(left.hash + right.hash), left=left, right=right)
                left.parent = parent
                right.parent = parent
                next_level.append(parent)
                i += 2
            queue = next_level
        return queue[0]

    def root_hash(self) -> bytes:
        return self.root.hash

    def proof_at(self, index: int) -> dict:
        """Return the Merkle proof (lemma, path) for a leaf index."""
        if index < 0 or index >= len(self.leaf_nodes):
            raise IndexError("index out of bound")
        if len(self.leaf_nodes) == 1:
            return {"lemma": [self.root.hash], "path": []}

        lemma: List[bytes] = [self.leaf_nodes[index].hash]
        path: List[bool] = []

        current = self.leaf_nodes[index]
        while current is not self.root:
            if current.is_left():
                lemma.append(current.parent.right.hash)
                path.append(True)
            else:
                lemma.append(current.parent.left.hash)
                path.append(False)
            current = current.parent

        lemma.append(self.root.hash)
        return {"lemma": lemma, "path": path}
