# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Merkle Tree

Provides Merkle tree construction, proof generation, and verification
for batching chain blocks into compact roots that can be anchored to L2.

Used by the chain_anchor module to batch multiple genesis chain blocks
into a single L2 transaction via Merkle root anchoring.

Dependencies: None (pure stdlib)
"""

import hashlib
import logging
from typing import Optional

logger = logging.getLogger(__name__)


class MerkleTree:
    """
    Binary Merkle tree for batching block hashes.

    Supports:
      - Construction from a list of hex hash strings
      - Inclusion proof generation for any leaf
      - Static proof verification
    """

    def __init__(self, hashes: list):
        """
        Build a Merkle tree from a list of hex hash strings.

        Args:
            hashes: List of hex-encoded hash strings (leaf nodes)
        """
        if not hashes:
            raise ValueError("Cannot build Merkle tree from empty list")

        self._leaves = list(hashes)
        self._layers = []
        self._root = self._build()

    @property
    def root(self) -> str:
        """The Merkle root hash."""
        return self._root

    @property
    def depth(self) -> int:
        """Tree depth (number of layers excluding leaves)."""
        return len(self._layers) - 1

    @property
    def leaf_count(self) -> int:
        """Number of leaf nodes."""
        return len(self._leaves)

    def _build(self) -> str:
        """Construct the Merkle tree and return the root hash."""
        layer = list(self._leaves)
        self._layers = [layer]

        while len(layer) > 1:
            next_layer = []
            for i in range(0, len(layer), 2):
                left = layer[i]
                # If odd number, duplicate the last element
                right = layer[i + 1] if i + 1 < len(layer) else left
                combined = hashlib.sha256((left + right).encode()).hexdigest()
                next_layer.append(combined)
            layer = next_layer
            self._layers.append(layer)

        return layer[0]

    def get_proof(self, index: int) -> list:
        """
        Generate an inclusion proof for the leaf at the given index.

        Args:
            index: Zero-based leaf index

        Returns:
            List of (sibling_hash, direction) tuples where direction is
            "left" or "right" indicating the sibling's position.
        """
        if index < 0 or index >= len(self._leaves):
            raise IndexError(f"Leaf index {index} out of range [0, {len(self._leaves)})")

        proof = []
        idx = index

        for layer in self._layers[:-1]:  # all layers except root
            if idx % 2 == 0:
                # Sibling is to the right
                sibling_idx = idx + 1
                direction = "right"
            else:
                # Sibling is to the left
                sibling_idx = idx - 1
                direction = "left"

            if sibling_idx < len(layer):
                proof.append((layer[sibling_idx], direction))
            else:
                # Odd layer, duplicate self
                proof.append((layer[idx], "right"))

            idx = idx // 2

        return proof

    @staticmethod
    def verify_proof(leaf: str, proof: list, root: str) -> bool:
        """
        Verify a Merkle inclusion proof.

        Args:
            leaf: The leaf hash to verify
            proof: List of (sibling_hash, direction) tuples
            root: Expected Merkle root

        Returns:
            True if the proof is valid
        """
        current = leaf

        for sibling, direction in proof:
            if direction == "left":
                combined = sibling + current
            else:
                combined = current + sibling
            current = hashlib.sha256(combined.encode()).hexdigest()

        return current == root

    def to_dict(self) -> dict:
        """Serialize the tree for storage."""
        return {
            "root": self._root,
            "leaves": self._leaves,
            "depth": self.depth,
            "leaf_count": self.leaf_count,
        }


def batch_chain_blocks(store, batch_size: int = 100) -> Optional[dict]:
    """
    Batch recent chain blocks into a Merkle root for L2 anchoring.

    Reads the last `batch_size` blocks from the ChainStore,
    builds a Merkle tree from their hashes, and returns the root
    with metadata.

    Args:
        store: ChainStore instance
        batch_size: Number of blocks to batch

    Returns:
        Dict with root, block range, depth, and leaf count.
        None if no blocks available.
    """
    chain_length = store.chain_length()
    if chain_length == 0:
        return None

    start = max(0, chain_length - batch_size)
    blocks = store.get_chain(start=start, limit=batch_size)

    if not blocks:
        return None

    hashes = [b.block_hash for b in blocks]
    tree = MerkleTree(hashes)

    return {
        "merkle_root": tree.root,
        "block_range": {
            "start": blocks[0].index,
            "end": blocks[-1].index,
        },
        "depth": tree.depth,
        "leaf_count": tree.leaf_count,
        "batched_at": __import__("datetime").datetime.now().isoformat(),
    }
