import hashlib
import json
from typing import List
from .crypto import Signer
import math

def _hash_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()

def merkle_root(digests: List[str]) -> str:
    if not digests:
        return ""
    level = [bytes.fromhex(d) for d in digests]
    while len(level) > 1:
        next_level = []
        for i in range(0, len(level), 2):
            left = level[i]
            right = level[i+1] if i+1 < len(level) else level[i]
            next_level.append(hashlib.sha256(left + right).digest())
        level = next_level
    return level[0].hex()

def merkle_path_for_index(digests: List[str], index: int) -> dict:
    """Compute Merkle root and the Merkle path for a given leaf index.

    Returns a dict with keys:
      - root: the Merkle root for the given leaves
      - path: list of sibling digests (hex) from leaf to root
    """
    if not digests:
        return {"root": "", "path": []}
    # convert hex digests to bytes for hashing
    level = [bytes.fromhex(d) for d in digests]
    idx = index
    path = []
    while len(level) > 1:
        # determine sibling index for current level
        sib_index = idx ^ 1  # toggle last bit -> sibling index
        if 0 <= sib_index < len(level):
            path.append(level[sib_index].hex())
        else:
            # no sibling, duplicate the current node
            path.append(level[idx].hex())
        # build next level
        next_level = []
        for i in range(0, len(level), 2):
            left = level[i]
            right = level[i+1] if i+1 < len(level) else level[i]
            next_level.append(hashlib.sha256(left + right).digest())
        level = next_level
        idx = idx // 2
    root = level[0].hex()
    return {"root": root, "path": path}

def verify_delta_root(delta_entries: List[dict]) -> bool:
    """Verify that a sequence of delta entries shares a consistent Merkle root.

    Each delta entry is expected to contain a top-level "digest" field, which
    is the SHA-256 digest of the serialized entry payload. The function computes
    the Merkle root over all provided digests and ensures that every delta entry
    that carries a "delta_root" field agrees with the computed root. If no
    delta_root is present, the function returns True once a root can be derived
    from the available digests.

    This helper is intended for cross-node verification when exchanging delta blocks.
    """
    if not delta_entries:
        return True

    # Collect all available digests in the delta entries
    digests = [e.get("digest") for e in delta_entries if e.get("digest") is not None]
    if not digests:
        return True

    computed_root = merkle_root(digests)
    # If entries provide an explicit delta_root, ensure consistency across them
    for e in delta_entries:
        explicit = e.get("delta_root")
        if explicit is not None and explicit != computed_root:
            return False
    return True


def verify_inclusion(leaf_digest: str, merkle_path: List[str], root: str, index: int) -> bool:
    """Verify that a leaf digest is included in a Merkle tree with given root.

    leaf_digest: hex string of the leaf digest
    merkle_path: list of sibling digests (hex) from leaf up to, but not including, the root
    root: expected Merkle root (hex)
    index: the zero-based index of the leaf within the original leaves array

    This mirrors the path generation in merkle_path_for_index: when a sibling is
    missing the leaf is duplicated. We reconstruct the hash up to the root and
    compare with the provided root.
    """
    if not leaf_digest:
        return False
    try:
        node = bytes.fromhex(leaf_digest)
    except Exception:
        return False

    idx = index
    for sib_hex in merkle_path:
        try:
            sib = bytes.fromhex(sib_hex)
        except Exception:
            return False
        if idx % 2 == 0:
            # current node is left child
            node = hashlib.sha256(node + sib).digest()
        else:
            # current node is right child
            node = hashlib.sha256(sib + node).digest()
        idx = idx // 2

    return node.hex() == (root or "")


class BloomFilter:
    """A minimal Bloom filter implementation used for compact existence indexes.

    This is intentionally small and dependency-free: it stores bits in a bytearray
    and uses multiple SHA-256-based hash functions with different salts.
    """
    def __init__(self, m: int = 1024, k: int = 3):
        if m <= 0:
            raise ValueError("m (number of bits) must be > 0")
        if k <= 0:
            raise ValueError("k (number of hash functions) must be > 0")
        self.m = int(m)
        self.k = int(k)
        self._bits = bytearray((self.m + 7) // 8)

    def _hashes(self, item: str):
        if isinstance(item, str):
            item = item.encode("utf-8")
        for i in range(self.k):
            h = hashlib.sha256()
            # include the index as a single byte salt
            h.update(i.to_bytes(1, "little"))
            h.update(item)
            yield int.from_bytes(h.digest(), "big") % self.m

    def add(self, item: str) -> None:
        for pos in self._hashes(item):
            byte_index = pos // 8
            bit_index = pos % 8
            self._bits[byte_index] |= (1 << bit_index)

    def __contains__(self, item: str) -> bool:
        for pos in self._hashes(item):
            byte_index = pos // 8
            bit_index = pos % 8
            if not (self._bits[byte_index] & (1 << bit_index)):
                return False
        return True

    def to_dict(self) -> dict:
        return {"m": self.m, "k": self.k, "bits": self._bits.hex()}

    @staticmethod
    def from_dict(d: dict) -> "BloomFilter":
        bf = BloomFilter(m=d.get("m", 1024), k=d.get("k", 3))
        bits_hex = d.get("bits", "")
        if bits_hex:
            bf._bits = bytearray.fromhex(bits_hex)
        return bf

class DeltaLog:
    def __init__(self):
        self.entries = []  # each entry is a dict with digest and payload
        self.anchor = None  # optional cloud/ground anchor for global verifiability
        # Checkpoints store archived roots and the set of digests included at that time
        # Each checkpoint is a dict: {"root": <hex>, "index": <exclusive_index>, "digests": [<hex>, ...]}
        self.checkpoints = []

    def add_entry(self, entry: dict) -> str:
        # entry must be serializable, and we store a digest for Merkle
        payload_bytes = json.dumps(entry, sort_keys=True).encode('utf-8')
        digest = hashlib.sha256(payload_bytes).hexdigest()
        self.entries.append({
            "digest": digest,
            "payload": entry,
        })
        return digest

    def anchor_root(self, anchor: str) -> None:
        """Record an optional anchor for the current delta log.

        This does not alter existing entries; it simply stores a reference
        to a trusted anchor (e.g., ground control) to tie the local log to
        an external verifiable state.
        """
        self.anchor = anchor

    def delta_from_index(self, index: int) -> List[dict]:
        # Return full LedgerEntry dictionaries augmented with Merkle proofs.
        # This enables recipients to reconstruct and verify the delta with
        # compact proofs, while preserving compatibility with existing code.
        leaves = [e["digest"] for e in self.entries]
        result = []
        for i in range(index, len(self.entries)):
            entry = self.entries[i]
            # Compute Merkle path for this leaf with respect to all leaves
            path_info = merkle_path_for_index(leaves, i)
            # The payload is the LedgerEntry dict previously stored in 'payload'
            full_entry = entry["payload"].copy()
            # Attach meta for verification without altering the payload structure
            full_entry.update({
                "digest": entry["digest"],
                "delta_root": path_info["root"],
                "merkle_path": path_info["path"],
            })
            result.append(full_entry)
        return result

    def root(self) -> str:
        digests = [e["digest"] for e in self.entries]
        return merkle_root(digests)

    def export_delta_block(self, since_index: int = 0, signer_key: bytes = None) -> dict:
        """Export a compact delta block suitable for DTN-friendly transfer.

        The block contains:
          - entries: full LedgerEntry dicts augmented with digest, delta_root, merkle_path
          - digests: list of leaf digests included in the block (in order)
          - root: merkle root for the block's leaves
          - anchor: optional anchor recorded on this log
          - aggregated_sig: HMAC-SHA256 of the root if signer_key is provided (hex)

        This is a lightweight, backwards-compatible representation that
        recipients can verify using verify_delta_root and (optional) aggregated_sig.
        """
        leaves = [e["digest"] for e in self.entries]
        # entries since index
        entries = []
        for i in range(since_index, len(self.entries)):
            entry = self.entries[i]
            path_info = merkle_path_for_index(leaves, i)
            full_entry = entry["payload"].copy()
            full_entry.update({
                "digest": entry["digest"],
                "delta_root": path_info["root"],
                "merkle_path": path_info["path"],
                "index": i,
            })
            entries.append(full_entry)

        block_digests = [e["digest"] for e in self.entries[since_index:]]
        block_root = merkle_root(block_digests)

        aggregated_sig = None
        if signer_key is not None:
            signer = Signer(signer_key)
            # sign the block root bytes
            aggregated_sig = signer.sign(block_root.encode("utf-8")).hex()

        return {
            "entries": entries,
            "digests": block_digests,
            "root": block_root,
            "anchor": self.anchor,
            "aggregated_sig": aggregated_sig,
            "since_index": since_index,
        }

    def checkpoint(self, prune_before_index: int = 0) -> dict:
        """Create a checkpoint for entries up to prune_before_index (exclusive) and optionally prune them.

        The checkpoint records the Merkle root over the archived leaves and the list of archived digests.
        It then removes those entries from the active entries list so storage can be reclaimed.

        Returns the checkpoint dict.
        """
        if prune_before_index <= 0:
            # nothing to archive
            cp_digests = []
            cp_root = ""
            cp_bloom = None
        else:
            cp_digests = [e["digest"] for e in self.entries[:prune_before_index]]
            cp_root = merkle_root(cp_digests)
            # Build a compact Bloom filter summarizing archived digests to
            # allow peers to quickly test probable presence before requesting
            # proofs or shards. Choose conservative defaults for m and k for
            # typical small checkpoints; these can be tuned per mission.
            m = max(1024, len(cp_digests) * 16)  # bits
            k = 4
            bf = BloomFilter(m=m, k=k)
            for d in cp_digests:
                bf.add(d)
            cp_bloom = bf.to_dict()

        checkpoint = {
            "root": cp_root,
            "index": prune_before_index,
            "digests": cp_digests,
            "bloom": cp_bloom,
        }
        # store checkpoint
        self.checkpoints.append(checkpoint)
        # prune entries up to prune_before_index
        if prune_before_index > 0:
            self.entries = self.entries[prune_before_index:]
        return checkpoint

    def has_digest(self, digest: str) -> bool:
        """Return True if digest is present in active entries or archived checkpoints."""
        for e in self.entries:
            if e.get("digest") == digest:
                return True
        for cp in self.checkpoints:
            if digest in cp.get("digests", []):
                return True
        return False


def verify_delta_block(block: dict, signer_key: bytes = None) -> bool:
    """Verify the integrity of a delta block.

    Checks performed:
      - recompute merkle root from provided digests and compare to block['root']
      - ensure each entry's digest matches the corresponding digest
      - if signer_key is provided, verify aggregated_sig matches HMAC of root
    """
    if not block:
        return False
    digests = block.get("digests", [])
    computed_root = merkle_root(digests)
    if computed_root != (block.get("root") or ""):
        return False

    # ensure entries line up with digests
    entries = block.get("entries", [])
    for i, entry in enumerate(entries):
        if entry.get("digest") != digests[i]:
            return False

    # verify optional aggregated signature
    if signer_key is not None:
        sig_hex = block.get("aggregated_sig")
        if sig_hex is None:
            return False
        signer = Signer(signer_key)
        expected = signer.sign((computed_root or "").encode("utf-8"))
        try:
            return expected.hex() == sig_hex
        except Exception:
            return False

    return True