"""A group's chat keys (epochs) and its stored history."""

from __future__ import annotations

import logging
import time as _time

from meshbay_common.chatbox import new_epoch_key
from meshbay_common.crypto import unwrap_gek_aes, wrap_gek_aes

from meshbay_node.ops.core import OpError, _group_ctx

log = logging.getLogger("meshbay_node.ops")


# ── Chat epoch keys ──────────────────────────────────────────────────────────
#
# The key a group's chat archive is encrypted under. Generated here, by the
# node, and never by a member — the C5b rule is about key material arriving from
# outside, and this is the same rule that lets `gek_rotate` be a signed
# instruction rather than a delivery.
#
# An *epoch* rather than a rotation, and the distinction is the whole design:
# opening a new one stops a departing member reading what comes next, while
# every earlier epoch is kept and still delivered to current members, so the
# history they could already read stays readable. Rotating instead — replacing
# the key, as `set_gek` does — would make every message anyone ever sent
# permanently unreadable to everybody, which is what a plain GEK-derived
# archive key would have done on the very first `member unpin`
# (finding F4, docs/MESHBAY_DESIGN.md §13.6).


async def _wrap_for_node(state: dict, key: bytes) -> dict:
    """
    Wrap a key to the node's own X25519 key, the way `set_gek` does for the GEK.

    Wrapped, not raw: the claim chat encryption makes is against someone who
    obtains the node's storage *without the keystore password*, and the node's
    X25519 private key is what the keystore protects. A raw key in SQLite would
    leave nothing behind that claim.
    """
    pk_x_node_raw = state.get("pk_x25519_raw")
    if not pk_x_node_raw:
        raise OpError("Node identity not available", status=503)
    return wrap_gek_aes(key, pk_x_node_raw)


async def chat_epoch_keys(state: dict, group_id: str) -> list[dict]:
    """
    Every chat epoch key this group has, oldest first, in the clear *in memory*.

    Cached on the group context: unwrapping is an ECIES operation per epoch and
    this is on the path of every member connecting to a group with chat on.
    """
    ctx = _group_ctx(state, group_id)
    cached = ctx.get("chat_epoch_keys")
    if cached is not None:
        return cached

    bundle_store = state.get("bundle_store")
    if not bundle_store:
        raise OpError("Bundle store not available", status=503)
    sk_x_raw = state.get("sk_x25519_raw")
    pk_x_raw = state.get("pk_x25519_raw")
    if not (sk_x_raw and pk_x_raw):
        raise OpError("Node identity not available", status=503)

    keys: list[dict] = []
    for row in await bundle_store.fetch_chat_epochs(group_id):
        try:
            keys.append({"epoch": row["epoch"],
                         "key": unwrap_gek_aes(row, sk_x_raw, pk_x_raw)})
        except Exception as e:
            # Loud, and not fatal: one unreadable epoch must not take the
            # readable ones with it. The messages of that epoch are lost, which
            # is a thing the operator needs told rather than a thing to hide.
            log.error("chat: epoch %d of group %s will not unwrap (%s) — "
                      "its messages are unreadable", row["epoch"],
                      group_id[:8], e)
    ctx["chat_epoch_keys"] = keys
    return keys


async def open_chat_epoch(state: dict, group_id: str) -> dict:
    """
    Open a new chat epoch. Idempotent only in the sense that it always adds one.

    Called when the set of devices that may read *future* messages shrinks: a
    member removed, a device revoked or unpinned, the group key rotated, or the
    operator asking directly. Never on a schedule — an epoch nobody needed is an
    epoch key the node has to keep for ever.
    """
    bundle_store = state.get("bundle_store")
    if not bundle_store:
        raise OpError("Bundle store not available", status=503)

    epoch = await bundle_store.latest_chat_epoch(group_id) + 1
    key = new_epoch_key()
    wrapped = await _wrap_for_node(state, key)
    await bundle_store.store_chat_epoch(
        group_id, epoch, wrapped["pk_eph_b64"], wrapped["nonce_b64"],
        wrapped["wrapped_b64"])

    # Tolerant of a group context that does not exist yet: the daemon opens the
    # first epoch **while it is building** `groups_ctx`, before publishing it on
    # the state, because a group with no epoch key is a group nobody can speak
    # in. Insisting on the context here would make start-up the one moment this
    # cannot be called.
    ctx = (state.get("groups_ctx") or {}).get(group_id)
    if ctx is not None:
        cached = ctx.get("chat_epoch_keys")
        if cached is not None:
            cached.append({"epoch": epoch, "key": key})
        ctx["chat_epoch"] = epoch

    # The transports hold their own view of the group, exactly as `set_gek`
    # notes: an epoch that did not reach them would have members sealing under
    # a key the node no longer thinks is current.
    for transport_key in ("webrtc", "quic_server"):
        transport = state.get(transport_key)
        groups = getattr(transport, "_ctx", {}).get("groups") if transport else None
        if groups and group_id in groups:
            groups[group_id]["chat_epoch"] = epoch
            groups[group_id].pop("chat_epoch_keys", None)

    log.info("Chat epoch %d opened for group %s", epoch, group_id[:8])
    return {"epoch": epoch}


async def ensure_chat_epoch(state: dict, group_id: str) -> int:
    """The current epoch, opening the first one if the group has none."""
    bundle_store = state.get("bundle_store")
    if not bundle_store:
        raise OpError("Bundle store not available", status=503)
    epoch = await bundle_store.latest_chat_epoch(group_id)
    if epoch:
        return epoch
    return (await open_chat_epoch(state, group_id))["epoch"]


async def chat_status(state: dict, group_id: str) -> dict:
    """What the operator needs to decide anything about this group's chat."""
    ctx = _group_ctx(state, group_id)
    bundle_store = state.get("bundle_store")
    store = ctx.get("chat_store")
    plain = sealed = 0
    if store is not None:
        plain, sealed = await store.count_by_format()
    return {
        "group_id": group_id,
        "epoch": (await bundle_store.latest_chat_epoch(group_id)
                  if bundle_store else 0),
        # Rows written before MNP 2.0. Not a state the node can be *in* — chat
        # is always encrypted now — but a state its disk can be in until
        # `chat encrypt-history` has run, and the operator has to be told,
        # because those messages are the ones still readable off a stolen disk.
        "plaintext_messages": plain,
        "encrypted_messages": sealed,
    }


async def encrypt_chat_history(state: dict, group_id: str) -> dict:
    """
    Re-encrypt the messages written before this group turned encryption on.

    Deliberately **not** done by the switch. It rewrites the only copy of a
    conversation, and a toggle that does that is one somebody flips twice; this
    is an explicit command, it copies the database first, and it runs in one
    transaction.

    The node can do this at all only because it holds those rows in plaintext —
    it is the last moment at which anyone can. Afterwards nothing on this
    machine can read them without an epoch key.

    Messages are sealed under a **synthetic device** belonging to the node, not
    under the original sender's key: the node does not hold anyone's signing key
    and must not pretend to. They are marked as such, so a reader is told these
    carry the node's word for who wrote them — which is all they ever carried,
    since they were written before signing existed.
    """
    import shutil

    from meshbay_common.chatbox import seal

    ctx = _group_ctx(state, group_id)
    store = ctx.get("chat_store")
    if store is None:
        raise OpError("This group has no chat store", status=404)

    epoch = await ensure_chat_epoch(state, group_id)
    keys = {k["epoch"]: k["key"] for k in await chat_epoch_keys(state, group_id)}
    key = keys.get(epoch)
    if not key:
        raise OpError("No chat key for this group", status=503)

    sk_node = state.get("sk_node")
    if sk_node is None:
        raise OpError("Node identity not available", status=503)
    from cryptography.hazmat.primitives import serialization

    device_raw = sk_node.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)
    import base64 as _b64

    device_b64 = _b64.b64encode(device_raw).decode()

    backup = store.db_path.with_name(
        f"{store.db_path.name}.bak-{int(_time.time())}")
    shutil.copy2(store.db_path, backup)

    converted = 0
    for row in await store.all_plaintext():
        text = (row.payload.decode("utf-8", errors="replace")
                if isinstance(row.payload, bytes) else str(row.payload))
        env = seal(key, group_id, epoch, device_b64, device_raw, sk_node, {
            "text": text,
            "thread_id": row.thread_id,
            "sender_name": row.sender_name,
            "sent_at": int(row.timestamp),
            # The node sealed this after the fact; it did not witness it being
            # signed. Said in the payload rather than inferred from the device.
            "migrated": True,
        })
        await store.reseal(row.id, epoch=epoch, device=device_raw,
                           nonce=env["nonce"], ct=env["ct"], sig=env["sig"])
        converted += 1
    await store.commit()

    log.info("Chat history re-encrypted for group %s: %d message(s), backup %s",
             group_id[:8], converted, backup.name)
    return {"group_id": group_id, "converted": converted,
            "backup": str(backup), "epoch": epoch}


async def prune_chat(state: dict, group_id: str, max_age_days: int) -> dict:
    """
    Delete messages older than `max_age_days`. Epoch keys are never touched.

    An epoch whose messages have all aged out costs 32 bytes and keeps the
    operation reversible in the only direction that matters: nothing that is
    still stored becomes unreadable because something else was deleted.
    """
    ctx = _group_ctx(state, group_id)
    store = ctx.get("chat_store")
    if store is None:
        raise OpError("This group has no chat store", status=404)
    if max_age_days < 1:
        raise OpError("max_age_days must be at least 1", status=400)
    removed = await store.delete_older_than(
        _time.time() - max_age_days * 86400)
    log.info("Chat retention for group %s: %d message(s) removed",
             group_id[:8], removed)
    return {"group_id": group_id, "removed": removed,
            "max_age_days": max_age_days}