"""
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
    _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)

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

    def serialize(self) -> bytes:
        """
        Produce a wire-ready byte string:
          msgpack → zstd → [GEK encrypt if private] → sign → length-prefixed envelope
        """
        payload = msgpack.packb({
            "group_id": self.group_id,
            "version":  self.version,
            "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"],
        )
        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."""
        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)

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