"""
Mesh Group Index — encrypted file listing for a group.

Wire format (private group):
  msgpack({entries, version, group_id}) → zstd compress → GEK ChaCha20 encrypt → sign

Wire format (public group):
  msgpack({entries, version, group_id}) → sign (no encryption)

Delta format:
  {base_version, version, additions: [...], deletions: [id, ...]}
"""

import base64
import logging
import os
import time
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Iterator

import blake3
import msgpack
import zstandard as zstd
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

from meshbay_common.crypto import (
    sign_chunk,
    verify_chunk_signature,
    pk_to_b64,
    generate_gek,
)
from meshbay_common.webcrypto import (
    chunk_key_aes as derive_chunk_key,
    encrypt_chunk_aes as encrypt_chunk,
    decrypt_chunk_aes as decrypt_chunk,
)
from meshbay_common.protocol import IndexEntry, IndexDelta

log = logging.getLogger(__name__)

ZSTD_LEVEL   = 3    # fast compression
INDEX_CHUNK  = 0    # the index itself is treated as chunk 0 of a virtual "index file"


@dataclass
class GroupIndex:
    """
    Encrypted, signed Mesh Group Index for one group.

    Usage:
        idx = GroupIndex(group_id="...", sk_node=sk, gek=gek_bytes)
        idx.add_entry(entry)
        wire_bytes = idx.serialize()          # for sending to members
        recovered  = GroupIndex.deserialize(wire_bytes, sk_node=sk, gek=gek_bytes)
    """

    group_id: str
    sk_node:  Ed25519PrivateKey
    gek:      bytes | None = None   # None → public group (no encryption)
    version:  int = 1
    # The group's roots and whether each is readable right now. Travels inside
    # the encrypted payload because it names the operator's directories, and a
    # member needs it to tell "temporarily unavailable" from "deleted" — a
    # distinction the entries alone cannot carry, since an unavailable root's
    # files are still listed. Absent in an index written before roots existed.
    roots:    list = field(default_factory=list)
    _entries: dict = field(default_factory=dict, repr=False)  # id → IndexEntry

    # ── Entry management ──────────────────────────────────────────────────────

    def add_entry(self, entry: IndexEntry) -> None:
        self._entries[entry.id] = entry

    def remove_entry(self, file_id: str) -> bool:
        return self._entries.pop(file_id, None) is not None

    def get_entry(self, file_id: str) -> IndexEntry | None:
        return self._entries.get(file_id)

    @property
    def entries(self) -> list[IndexEntry]:
        return list(self._entries.values())

    @property
    def count(self) -> int:
        return len(self._entries)

    def entries_by_id(self) -> dict:
        """A snapshot copy, for diff() to compare a later version against —
        see daemon.py _on_index_change, the only caller."""
        return dict(self._entries)

    @classmethod
    def _snapshot(cls, group_id: str, sk_node: Ed25519PrivateKey, gek: bytes | None,
                  version: int, entries_by_id: dict) -> "GroupIndex":
        """
        A lightweight stand-in for diff()'s `previous` argument — never
        serialized or sent anywhere, just a comparison point built from an
        earlier entries_by_id() snapshot rather than a live GroupIndex.
        """
        idx = cls(group_id=group_id, sk_node=sk_node, gek=gek, version=version)
        idx._entries = dict(entries_by_id)
        return idx

    # ── Serialisation ─────────────────────────────────────────────────────────

    def serialize(self) -> bytes:
        """
        Produce a signed index envelope:
          msgpack → zstd → [GEK encrypt if private] → sign → length-prefixed envelope

        **This is not an MNP message.** It was the payload of `index_sync` on the QUIC
        transport, while WebRTC sent plain entries under the same type — one message
        type, two encodings (2026-09-03). Both transports now build `index_sync` from
        `transport/wire.py`. This stays as a correct at-rest/interchange format, and
        as the only thing that signs and encrypts a whole index; read it as that, not
        as a wire contract.
        """
        payload = msgpack.packb({
            "group_id": self.group_id,
            "version":  self.version,
            "roots":    list(self.roots),
            "entries":  [asdict(e) for e in self.entries],
        }, use_bin_type=True)

        compressed = zstd.compress(payload, level=ZSTD_LEVEL)

        if self.gek is not None:
            # Private group: encrypt with GEK-derived key
            idx_hash   = blake3.blake3(compressed).digest()
            ckey       = derive_chunk_key(self.gek, idx_hash, INDEX_CHUNK)
            nonce, ct  = encrypt_chunk(ckey, compressed)
            sig        = sign_chunk(self.sk_node, INDEX_CHUNK, nonce, blake3.blake3(ct).digest())
            envelope = msgpack.packb({
                "type":       "index",
                "encrypted":  True,
                "version":    self.version,
                "group_id":   self.group_id,
                "idx_hash_b64": base64.b64encode(idx_hash).decode(),
                "nonce_b64":  base64.b64encode(nonce).decode(),
                "ct_b64":     base64.b64encode(ct).decode(),
                "sig_b64":    base64.b64encode(sig).decode(),
                "pk_node_b64": pk_to_b64(self.sk_node.public_key()),
            }, use_bin_type=True)
        else:
            # Public group: just sign the compressed payload
            payload_hash = blake3.blake3(compressed).digest()
            sig = sign_chunk(self.sk_node, INDEX_CHUNK,
                             bytes(12),        # zero nonce for plaintext
                             payload_hash)
            envelope = msgpack.packb({
                "type":       "index",
                "encrypted":  False,
                "version":    self.version,
                "group_id":   self.group_id,
                "data_b64":   base64.b64encode(compressed).decode(),
                "hash_b64":   base64.b64encode(payload_hash).decode(),
                "sig_b64":    base64.b64encode(sig).decode(),
                "pk_node_b64": pk_to_b64(self.sk_node.public_key()),
            }, use_bin_type=True)

        return envelope

    @classmethod
    def deserialize(
        cls,
        data: bytes,
        sk_node: Ed25519PrivateKey,
        gek: bytes | None = None,
    ) -> "GroupIndex":
        """Deserialize, verify signature, and decrypt (if private)."""
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
        from meshbay_common.crypto import verify_chunk_signature

        envelope = msgpack.unpackb(data, raw=False)

        pk_node_raw = base64.b64decode(envelope["pk_node_b64"])
        pk_node     = Ed25519PublicKey.from_public_bytes(pk_node_raw)
        sig         = base64.b64decode(envelope["sig_b64"])

        if envelope["encrypted"]:
            if gek is None:
                raise ValueError("GEK required to decrypt private group index")
            ct        = base64.b64decode(envelope["ct_b64"])
            nonce     = base64.b64decode(envelope["nonce_b64"])
            ct_hash   = blake3.blake3(ct).digest()
            verify_chunk_signature(pk_node, INDEX_CHUNK, nonce, ct_hash, sig)

            idx_hash  = base64.b64decode(envelope["idx_hash_b64"])
            ckey      = derive_chunk_key(gek, idx_hash, INDEX_CHUNK)
            compressed = decrypt_chunk(ckey, nonce, ct)
        else:
            compressed  = base64.b64decode(envelope["data_b64"])
            payload_hash = base64.b64decode(envelope["hash_b64"])
            verify_chunk_signature(pk_node, INDEX_CHUNK, bytes(12), payload_hash, sig)

        payload = msgpack.unpackb(zstd.decompress(compressed), raw=False)
        idx = cls(
            group_id=payload["group_id"],
            sk_node=sk_node,
            gek=gek,
            version=payload["version"],
            # Absent from an index written before roots existed; an empty list
            # reads as "nothing known about availability", not "no roots".
            roots=payload.get("roots") or [],
        )
        for e in payload["entries"]:
            idx.add_entry(IndexEntry(**e))
        return idx

    # ── Delta ─────────────────────────────────────────────────────────────────

    def diff(self, previous: "GroupIndex") -> IndexDelta:
        """
        Compute what changed since a previous version of this index.

        A shared id whose entry object now compares unequal (field-by-field,
        via IndexEntry's dataclass-generated __eq__) is an update, not an
        addition — the Videos app's async enrichment (duration, thumb_hash,
        title, ...) replaces an existing entry's fields after the fact via
        `add_entry`, which never introduces a new id. This only works
        because that replacement always constructs a *new* IndexEntry object
        (`dataclasses.replace`, never in-place attribute mutation) — mutating
        the same object in place would also mutate `previous`'s copy, since
        entries_by_id() is a shallow dict copy, and the two would always
        compare equal.
        """
        prev_ids = set(previous._entries)
        curr_ids = set(self._entries)

        additions = [self._entries[i] for i in curr_ids - prev_ids]
        deletions = list(prev_ids - curr_ids)
        updates = [
            self._entries[i] for i in curr_ids & prev_ids
            if self._entries[i] != previous._entries[i]
        ]

        return IndexDelta(
            base_version=previous.version,
            version=self.version,
            additions=additions,
            deletions=deletions,
            updates=updates,
        )