"""
Merkle tree and append-only log for Verisoo.

Inspired by Certificate Transparency (RFC 6962). Each attestation record
becomes a leaf in a merkle tree. The tree structure allows:

- Efficient proof that a record exists (O(log n) proof size)
- Efficient comparison between nodes ("do we agree on history?")
- Tamper evidence (any modification changes the root hash)

The log is append-only: records can be added but never removed or modified.
This gives Byzantine fault tolerance - if one node goes rogue, others have
the canonical record.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterator

from .exceptions import MerkleError
from .models import AttestationRecord


@dataclass
class MerkleNode:
    """A node in the merkle tree."""

    hash: str
    left: MerkleNode | None = None
    right: MerkleNode | None = None
    leaf_data: bytes | None = None  # Only for leaf nodes

    @property
    def is_leaf(self) -> bool:
        return self.leaf_data is not None


@dataclass
class InclusionProof:
    """
    Proof that a record exists in the tree.

    Contains the sibling hashes needed to recompute the root.
    Verifier can check: hash(record) + proof_hashes = known_root
    """

    leaf_hash: str
    leaf_index: int
    tree_size: int
    proof_hashes: list[str]
    root_hash: str


@dataclass
class ConsistencyProof:
    """
    Proof that an older tree is a prefix of a newer tree.

    Used during sync: "I have root X with N entries, you have root Y with M entries.
    Prove that your tree includes all my entries."
    """

    old_size: int
    new_size: int
    proof_hashes: list[str]


class MerkleLog:
    """
    Append-only merkle log for attestation records.

    Thread-safe for reads, single-writer for appends.
    Persistence is handled separately (see storage module).
    """

    def __init__(self) -> None:
        self._leaves: list[bytes] = []  # Serialized records
        self._leaf_hashes: list[str] = []
        self._root: str | None = None

    @property
    def size(self) -> int:
        """Number of records in the log."""
        return len(self._leaves)

    @property
    def root_hash(self) -> str | None:
        """Current merkle root, or None if empty."""
        return self._root

    def append(self, record: AttestationRecord) -> int:
        """
        Append a record to the log.

        Returns the index of the new record.
        """
        record_bytes = self._serialize_record(record)
        leaf_hash = self._hash_leaf(record_bytes)

        self._leaves.append(record_bytes)
        self._leaf_hashes.append(leaf_hash)
        self._root = self._compute_root()

        return len(self._leaves) - 1

    def get(self, index: int) -> AttestationRecord:
        """Get record at index."""
        if index < 0 or index >= len(self._leaves):
            raise MerkleError(f"Index {index} out of range [0, {len(self._leaves)})")
        return self._deserialize_record(self._leaves[index])

    def inclusion_proof(self, index: int) -> InclusionProof:
        """
        Generate proof that record at index exists in the tree.

        The proof is O(log n) hashes that let anyone verify
        the record is part of the tree with the given root.
        """
        if index < 0 or index >= len(self._leaves):
            raise MerkleError(f"Index {index} out of range")

        proof_hashes = self._build_inclusion_proof(index, len(self._leaves))

        return InclusionProof(
            leaf_hash=self._leaf_hashes[index],
            leaf_index=index,
            tree_size=len(self._leaves),
            proof_hashes=proof_hashes,
            root_hash=self._root or "",
        )

    def consistency_proof(self, old_size: int) -> ConsistencyProof:
        """
        Generate proof that tree at old_size is prefix of current tree.

        Used during federation: prove we haven't rewritten history.
        """
        if old_size < 0 or old_size > len(self._leaves):
            raise MerkleError(f"Old size {old_size} invalid for tree of size {len(self._leaves)}")

        proof_hashes = self._build_consistency_proof(old_size, len(self._leaves))

        return ConsistencyProof(
            old_size=old_size,
            new_size=len(self._leaves),
            proof_hashes=proof_hashes,
        )

    def verify_inclusion(self, proof: InclusionProof) -> bool:
        """Verify an inclusion proof against our current root."""
        if self._root is None:
            return False

        # Recompute root from leaf + proof
        computed = proof.leaf_hash
        index = proof.leaf_index
        size = proof.tree_size

        for sibling_hash in proof.proof_hashes:
            if index % 2 == 0:
                # We're left child, sibling is right
                computed = self._hash_pair(computed, sibling_hash)
            else:
                # We're right child, sibling is left
                computed = self._hash_pair(sibling_hash, computed)
            index //= 2

        return computed == proof.root_hash

    def __iter__(self) -> Iterator[AttestationRecord]:
        """Iterate over all records in order."""
        for record_bytes in self._leaves:
            yield self._deserialize_record(record_bytes)

    def __len__(self) -> int:
        return self.size

    # --- Internal methods ---

    def _hash_leaf(self, data: bytes) -> str:
        """Hash a leaf node (prefixed with 0x00 to distinguish from internal nodes)."""
        return hashlib.sha256(b"\x00" + data).hexdigest()

    def _hash_pair(self, left: str, right: str) -> str:
        """Hash two child nodes (prefixed with 0x01)."""
        combined = bytes.fromhex(left) + bytes.fromhex(right)
        return hashlib.sha256(b"\x01" + combined).hexdigest()

    def _compute_root(self) -> str:
        """Compute merkle root from current leaves."""
        if not self._leaf_hashes:
            return hashlib.sha256(b"").hexdigest()

        # Build tree bottom-up
        level = self._leaf_hashes.copy()

        while len(level) > 1:
            next_level = []
            for i in range(0, len(level), 2):
                if i + 1 < len(level):
                    next_level.append(self._hash_pair(level[i], level[i + 1]))
                else:
                    # Odd node: promote to next level
                    next_level.append(level[i])
            level = next_level

        return level[0]

    def _build_inclusion_proof(self, index: int, size: int) -> list[str]:
        """Build the sibling hashes for inclusion proof."""
        proof = []
        level = self._leaf_hashes.copy()

        while len(level) > 1:
            # Find sibling
            if index % 2 == 0:
                if index + 1 < len(level):
                    proof.append(level[index + 1])
            else:
                proof.append(level[index - 1])

            # Move to parent level
            next_level = []
            for i in range(0, len(level), 2):
                if i + 1 < len(level):
                    next_level.append(self._hash_pair(level[i], level[i + 1]))
                else:
                    next_level.append(level[i])
            level = next_level
            index //= 2

        return proof

    def _build_consistency_proof(self, old_size: int, new_size: int) -> list[str]:
        """Build consistency proof hashes."""
        # Simplified: return subtree roots that prove consistency
        # Full implementation would follow RFC 6962 algorithm
        # For now, return empty - federation will implement full version
        return []

    def _serialize_record(self, record: AttestationRecord) -> bytes:
        """Serialize record to bytes for storage/hashing."""
        # Use JSON for now - could use protobuf for efficiency
        data = {
            "sha256": record.image_hashes.sha256,
            "phash": record.image_hashes.phash,
            "dhash": record.image_hashes.dhash,
            "signature": record.signature.hex(),
            "attestor": record.attestor_fingerprint,
            "timestamp": record.timestamp.isoformat(),
            "metadata": record.metadata,
        }
        return json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8")

    def _deserialize_record(self, data: bytes) -> AttestationRecord:
        """Deserialize record from bytes."""
        from .models import ImageHashes

        obj = json.loads(data.decode("utf-8"))
        return AttestationRecord(
            image_hashes=ImageHashes(
                sha256=obj["sha256"],
                phash=obj["phash"],
                dhash=obj["dhash"],
            ),
            signature=bytes.fromhex(obj["signature"]),
            attestor_fingerprint=obj["attestor"],
            timestamp=datetime.fromisoformat(obj["timestamp"]),
            metadata=obj.get("metadata", {}),
        )


def verify_inclusion_proof(
    record_bytes: bytes,
    proof: InclusionProof,
    expected_root: str,
) -> bool:
    """
    Standalone verification of inclusion proof.

    Anyone can verify without having the full log.
    """
    leaf_hash = hashlib.sha256(b"\x00" + record_bytes).hexdigest()
    if leaf_hash != proof.leaf_hash:
        return False

    computed = leaf_hash
    index = proof.leaf_index

    for sibling_hash in proof.proof_hashes:
        combined = (
            bytes.fromhex(computed) + bytes.fromhex(sibling_hash)
            if index % 2 == 0
            else bytes.fromhex(sibling_hash) + bytes.fromhex(computed)
        )
        computed = hashlib.sha256(b"\x01" + combined).hexdigest()
        index //= 2

    return computed == expected_root
