"""
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
from dataclasses import asdict, dataclass, field

import blake3
import msgpack
import zstandard as zstd
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from meshbay_common.crypto import (
    pk_to_b64,
    sign_chunk,
    verify_chunk_signature,
)
from meshbay_common.protocol import IndexDelta, IndexEntry
from meshbay_common.webcrypto import (
    chunk_key_aes as derive_chunk_key,
)
from meshbay_common.webcrypto import (
    decrypt_chunk_aes as decrypt_chunk,
)
from meshbay_common.webcrypto import (
    encrypt_chunk_aes as encrypt_chunk,
)

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

        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,
        )