"""Binary SHA-256 Merkle tree over agent chain heads.

Provides global integrity verification: every agent's chain head is included
in a single signed Merkle root. Any agent can get an inclusion proof to verify
their chain is part of the global state.
"""

from __future__ import annotations

import hashlib
import math
from dataclasses import dataclass, field
from typing import Optional


def _sha256(data: bytes) -> str:
    """SHA-256 hash as hex string."""
    return hashlib.sha256(data).hexdigest()


def _hash_pair(left: str, right: str) -> str:
    """Hash two hex strings together: SHA-256(left || right)."""
    return _sha256((left + right).encode("utf-8"))


@dataclass
class MerkleTree:
    """Binary SHA-256 Merkle tree."""
    root: str  # hex string
    leaves: list[str]  # ordered leaf hashes
    leaf_labels: list[str]  # corresponding agent pubkeys
    levels: list[list[str]] = field(default_factory=list)  # all tree levels, bottom to top

    @property
    def leaf_count(self) -> int:
        return len(self.leaves)


def build_merkle_tree(chain_heads: dict[str, str]) -> MerkleTree:
    """Build a Merkle tree from agent chain heads.

    Args:
        chain_heads: {agent_pubkey: chain_head_hash} mapping.
                     Must be non-empty.

    Returns:
        MerkleTree with root, leaves, and internal structure.
    """
    if not chain_heads:
        empty_root = _sha256(b"empty")
        return MerkleTree(root=empty_root, leaves=[], leaf_labels=[], levels=[[empty_root]])

    # WHY sort by pubkey: Merkle root must be deterministic regardless of
    # insertion order. Sorting by pubkey (a hex string) gives a canonical
    # leaf ordering that any verifier can independently reconstruct.
    sorted_keys = sorted(chain_heads.keys())
    leaves = [chain_heads[k] for k in sorted_keys]

    # Build tree bottom-up
    levels = [leaves[:]]  # level 0 = leaves
    current = leaves[:]

    while len(current) > 1:
        next_level = []
        for i in range(0, len(current), 2):
            if i + 1 < len(current):
                next_level.append(_hash_pair(current[i], current[i + 1]))
            else:
                # WHY hash-with-self for odd leaves: Promotes the unpaired leaf
                # to the next level without special-casing. This is the standard
                # Bitcoin/Certificate Transparency approach. The alternative
                # (just promoting the hash unchanged) creates ambiguity in
                # inclusion proofs.
                next_level.append(_hash_pair(current[i], current[i]))
        levels.append(next_level)
        current = next_level

    root = current[0] if current else _sha256(b"empty")

    return MerkleTree(
        root=root,
        leaves=leaves,
        leaf_labels=sorted_keys,
        levels=levels,
    )


def get_inclusion_proof(tree: MerkleTree, agent_pubkey: str) -> Optional[list[dict]]:
    """Get an inclusion proof for an agent's chain head.

    Args:
        tree: The Merkle tree.
        agent_pubkey: The agent's public key.

    Returns:
        List of proof steps: [{"hash": str, "side": "left"|"right"}, ...]
        or None if the agent is not in the tree.
    """
    if agent_pubkey not in tree.leaf_labels:
        return None

    idx = tree.leaf_labels.index(agent_pubkey)
    proof = []

    for level in tree.levels[:-1]:  # skip root level
        if idx % 2 == 0:
            # We're the left sibling, need the right
            sibling_idx = idx + 1
            if sibling_idx < len(level):
                proof.append({"hash": level[sibling_idx], "side": "right"})
            else:
                # Odd-one-out, paired with itself
                proof.append({"hash": level[idx], "side": "right"})
        else:
            # We're the right sibling, need the left
            proof.append({"hash": level[idx - 1], "side": "left"})

        idx = idx // 2

    return proof


def verify_inclusion(root: str, leaf_hash: str, proof: list[dict]) -> bool:
    """Verify that a leaf is included in the Merkle tree.

    Args:
        root: The expected Merkle root.
        leaf_hash: The leaf hash to verify.
        proof: The inclusion proof from get_inclusion_proof().

    Returns:
        True if the proof is valid and the leaf is in the tree.
    """
    current = leaf_hash

    for step in proof:
        if step["side"] == "right":
            current = _hash_pair(current, step["hash"])
        else:
            current = _hash_pair(step["hash"], current)

    return current == root
