# SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 RE::AI Strategy
"""HDK Hash Engine — hierarchical SHA-256 hashing with genealogy."""
from __future__ import annotations

import hashlib
import hmac
import json
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from hdk.models import HDKNode


class HashEngine:
    """Produces hierarchical hashes where every hash encodes its full ancestry."""

    def __init__(self, algorithm: str = "sha256", node_key: bytes | None = None):
        """Initialize with hash algorithm (currently only sha256).

        `node_key` (0.1.2): a key derived from the canary secret (CanaryLayer.node_key). With it,
        the node_hash of every canary-woven node is an HMAC under that key, so a node cannot be
        rewritten and re-hashed by someone who does not hold the canary key.
        """
        if algorithm != "sha256":
            raise ValueError(f"Unsupported algorithm: {algorithm}. Only 'sha256' is supported.")
        self._algorithm = algorithm
        self.node_key = node_key

    def hash_content(self, content: bytes | str) -> str:
        """Hash raw content. Returns hex digest."""
        if isinstance(content, str):
            content = content.encode("utf-8")
        return hashlib.sha256(content).hexdigest()

    def hash_with_parents(self, content_hash: str, parent_chain: list[str]) -> str:
        """
        Hierarchical hash: H(content_hash || parent1 || parent2 || ... || root).

        This is the core HDK property — changing any ancestor
        invalidates this hash and all descendants.

        Args:
            content_hash: SHA-256 of the actual content
            parent_chain: [immediate_parent_hash, grandparent_hash, ..., root_hash]

        Returns:
            Hierarchical hash as hex string
        """
        combined = "||".join([content_hash] + parent_chain)
        return hashlib.sha256(combined.encode("utf-8")).hexdigest()

    # RFC 6962 domain separation: a leaf hash can never be confused with an interior node
    LEAF_PREFIX = b"\x00"
    NODE_PREFIX = b"\x01"

    @classmethod
    def merkle_leaf(cls, data: bytes | str) -> bytes:
        """Leaf hash: H(0x00 || data)."""
        if isinstance(data, str):
            data = data.encode("utf-8")
        return hashlib.sha256(cls.LEAF_PREFIX + data).digest()

    @classmethod
    def merkle_node(cls, left: bytes, right: bytes) -> bytes:
        """Interior node hash: H(0x01 || left || right)."""
        return hashlib.sha256(cls.NODE_PREFIX + left + right).digest()

    def compute_merkle_root(self, node_hashes: list[str]) -> str:
        """
        Merkle Tree Hash (RFC 6962, section 2.1) over a list of node hashes.

        leaf = H(0x00 || node_hash), node = H(0x01 || left || right). Leaves are never
        duplicated; for a level with an odd count the last *hashed* entry moves up a level,
        which yields the same root as RFC 6962's split at the largest power of two. A single
        leaf is still hashed, so the root never equals a bare input. Empty list: H("").

        Args:
            node_hashes: list of hash strings (each leaf's data is its UTF-8 encoding)

        Returns:
            Merkle root as hex string
        """
        if not node_hashes:
            return hashlib.sha256(b"").hexdigest()

        level = [self.merkle_leaf(h) for h in node_hashes]
        while len(level) > 1:
            nxt = [self.merkle_node(level[i], level[i + 1]) for i in range(0, len(level) - 1, 2)]
            if len(level) % 2:
                nxt.append(level[-1])
            level = nxt
        return level[0].hex()

    NODE_HASH_DOMAIN = b"hdk-node-v2\x00"
    #: 0.1.2: the keyed node hash (HMAC-SHA256 under the canary-derived node key)
    NODE_HMAC_DOMAIN = b"hdk-node-v3-hmac\x00"

    @staticmethod
    def canonical_json(obj) -> bytes:
        """Canonical serialisation: sorted keys, no whitespace, UTF-8."""
        return json.dumps(
            obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str
        ).encode("utf-8")

    def node_hash_fields(self, node: HDKNode) -> dict:
        """Every stored field of a node that node_hash commits to (all but node_hash itself)."""
        return {
            "id": node.id,
            "level": node.level.value if node.level else "",
            "timestamp": node.timestamp.isoformat() if node.timestamp else "",
            "content_hash": node.content_hash,
            "parent_chain": list(node.parent_chain),
            "self_hash": node.self_hash,
            "canary_commitment": node.canary_commitment,
            "canary_generation": node.canary_generation,
            "model": node.model or "",
            "provider": node.provider or "",
            "event_type": node.event_type.value if node.event_type else "",
            "token_count": node.token_count,
            "latency_ms": node.latency_ms,
            "user_id_token": node.user_id_token or "",
            "llm_model": node.llm_model or "",
            "llm_provider": node.llm_provider or "",
            "agent_name": node.agent_name or "",
            "custom": node.custom or {},
        }

    def create_node_hash(self, node: HDKNode) -> str:
        """
        Compute final node_hash for an HDKNode.

        node_hash = H("hdk-node-v2" || 0x00 || canonical_json(fields))
        or, with a node key and a canary-woven node (0.1.2):
        node_hash = HMAC-SHA256(node_key, "hdk-node-v3-hmac" || 0x00 || canonical_json(fields))

        where fields are all stored node fields (see node_hash_fields): the content and
        parent chain (via self_hash, and again explicitly), canary commitment, metadata,
        actor fields and custom fields. Canonical JSON makes the encoding unambiguous, so no
        two different field sets can serialise to the same bytes.
        """
        payload = self.canonical_json(self.node_hash_fields(node))
        if self.keyed(node):
            # 0.1.2: bound to the canary key — without it, a rewritten node cannot be re-hashed
            mac = hmac.new(self.node_key, self.NODE_HMAC_DOMAIN + payload, hashlib.sha256)
            return mac.hexdigest()
        return hashlib.sha256(self.NODE_HASH_DOMAIN + payload).hexdigest()

    def keyed(self, node: HDKNode) -> bool:
        """Whether `node` is hashed with the node key (a key is set and the node is woven)."""
        return self.node_key is not None and node.canary_generation > 0

    def verify_chain(self, node: HDKNode, parent_nodes: list[HDKNode]) -> bool:
        """
        Verify that a node's hash is consistent with its claimed parent chain.
        Recomputes self_hash, and checks that parent_nodes[i].node_hash == parent_chain[i]
        and that each parent's node_hash is itself genuine.

        Returns True if valid, False if any hash doesn't match.
        """
        # Recompute self_hash from content_hash + parent_chain
        expected_self_hash = self.hash_with_parents(node.content_hash, node.parent_chain)
        if expected_self_hash != node.self_hash:
            return False
        if node.node_hash and self.create_node_hash(node) != node.node_hash:
            return False

        # parent_chain stores each parent's node_hash (see HDK._create_node), so compare
        # against that; every claimed parent must be supplied and match.
        if len(parent_nodes) < len(node.parent_chain):
            return False
        for parent_hash, parent in zip(node.parent_chain, parent_nodes):
            if parent.node_hash != parent_hash:
                return False
            if self.create_node_hash(parent) != parent.node_hash:
                return False

        return True