"""
Chat encryption: what the node stores, what it refuses, and what survives.

Design A of `docs/MESHBAY_DESIGN.md` §4.5. Every test here is written as "this
does not work" or "this still works after X" — the regressions the plan's
register names, in the order they would bite.

The load-bearing ones are the last three. Rotation is the failure the design
exists to avoid: a chat key derived from the group key would have made every
message ever sent unreadable on the first `member unpin`, for everybody,
including the operator, and that is the *documented* procedure after removing
someone. Key storage is the failure that would make the whole feature a
decoration. Downgrade is C6's lesson, one feature later.
"""

import base64
import time
from pathlib import Path

import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from meshbay_common.chatbox import open_message, seal
from meshbay_common.crypto import generate_gek
from meshbay_common.groupbox import PURPOSE_CHAT_KEYS, unseal
from meshbay_common.protocol import MNP
from meshbay_node import ops
from meshbay_node.bundle_store import BundleStore
from meshbay_node.chat import FORMAT_SEALED_V1, ChatStore
from meshbay_node.indexer.group_index import GroupIndex
from meshbay_node.roster import open_roster
from meshbay_node.transport.webrtc_server import WebRTCPeerSession

from conftest import one_root

GROUP = "g" * 32


def _device():
    sk = Ed25519PrivateKey.generate()
    raw = sk.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)
    return sk, raw, base64.b64encode(raw).decode()


@pytest.fixture
async def node(tmp_path):
    """A daemon state with the pieces the chat path actually touches."""
    roster = await open_roster(tmp_path)
    bundles = BundleStore(tmp_path / "bundles.db")
    await bundles.open()
    chat = ChatStore(tmp_path / "chat.db")
    await chat.open()

    sk_x = Ed25519PrivateKey.generate()   # stand-in shape; X25519 below
    from cryptography.hazmat.primitives.asymmetric.x25519 import (
        X25519PrivateKey,
    )
    sk_x = X25519PrivateKey.generate()
    sk_x_raw = sk_x.private_bytes(
        serialization.Encoding.Raw, serialization.PrivateFormat.Raw,
        serialization.NoEncryption())
    pk_x_raw = sk_x.public_key().public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)

    shared = tmp_path / "shared"
    shared.mkdir(exist_ok=True)
    index = GroupIndex(group_id=GROUP, sk_node=Ed25519PrivateKey.generate())
    gek = generate_gek()

    group_ctx = {
        "gek": gek, "index": index, "roots": one_root(shared),
        "chat_store": chat, "chat_epoch": 0,
        "_peers": {},
    }
    state = {
        "roster": roster, "bundle_store": bundles,
        "sk_x25519_raw": sk_x_raw, "pk_x25519_raw": pk_x_raw,
        "groups_ctx": {GROUP: group_ctx}, "node_user_id": "operator",
    }
    yield {"state": state, "group_ctx": group_ctx, "gek": gek,
           "chat": chat, "roster": roster, "bundles": bundles,
           "index": index, "tmp_path": tmp_path}
    await chat.close()
    await bundles.close()
    await roster.close()


def _session(node, user_id="alice", device_b64=""):
    session = WebRTCPeerSession.__new__(WebRTCPeerSession)
    session._ctx = {"groups": {GROUP: node["group_ctx"]},
                    "daemon_state": node["state"]}
    session._group_id = GROUP
    session._user_id = user_id
    session._username = user_id
    session._pinned_pk = device_b64
    session._device_confirmed = bool(device_b64)
    session._registry_key = f"conn-{user_id}-{len(node['group_ctx']['_peers'])}"
    session.sent = []
    session._send = session.sent.append
    session._audit = lambda *a, **k: None
    return session


async def _drain(session, coro_holder):
    """`_spawn` stubbed to await inline, so a test sees the store written."""
    pass


def _spawn_inline(session):
    import asyncio

    pending = []
    session._spawn = lambda coro: pending.append(
        asyncio.get_event_loop().create_task(coro))
    return pending


async def _send_sealed(node, session, sk, device_raw, device_b64, text,
                       epoch=None):
    keys = await ops.chat_epoch_keys(node["state"], GROUP)
    epoch = epoch or keys[-1]["epoch"]
    key = next(k["key"] for k in keys if k["epoch"] == epoch)
    env = seal(key, GROUP, epoch, device_b64, device_raw, sk,
               {"text": text, "sender_name": session._user_id})
    pending = _spawn_inline(session)
    session._do_chat_message({
        "format": FORMAT_SEALED_V1, "epoch": epoch, "device": device_raw,
        "ct": env["ct"], "nonce": env["nonce"], "sig": env["sig"],
    })
    for task in pending:
        await task
    return env


# ── the archive survives what would destroy it ───────────────────────────────

async def test_history_survives_a_group_key_rotation(node):
    """
    R1, and the reason Design A exists.

    A chat key derived from the group key would be gone the moment the operator
    rotates — which is the documented step after removing a member. Every
    message ever sent would become unreadable, for everybody. The epoch key is
    wrapped under the group key *at delivery* and never stored under it, so a
    rotation is a re-wrap and costs nothing.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    sk, raw, b64 = _device()
    session = _session(node, device_b64=b64)
    await _send_sealed(node, session, sk, raw, b64, "before the rotation")

    # Rotate the group key, exactly as the operator does after a removal.
    node["group_ctx"]["gek"] = generate_gek()

    keys = await ops.chat_epoch_keys(node["state"], GROUP)
    stored = (await node["chat"].get_recent(10))[0]
    opened = open_message(
        next(k["key"] for k in keys if k["epoch"] == stored.epoch),
        GROUP, stored.epoch, b64, stored.nonce, stored.payload)
    assert opened["text"] == "before the rotation", (
        "rotating the group key must not make the chat archive unreadable — "
        "F4, and the whole reason the epoch key is not derived from it")


async def test_a_new_epoch_does_not_orphan_the_old_ones(node):
    """
    R2. Opening an epoch stops a removed member reading what comes *next*; it
    must leave what they could already read readable to everybody else.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    sk, raw, b64 = _device()
    session = _session(node, device_b64=b64)
    await _send_sealed(node, session, sk, raw, b64, "epoch one")

    await ops.open_chat_epoch(node["state"], GROUP)
    await _send_sealed(node, session, sk, raw, b64, "epoch two")

    keys = {k["epoch"]: k["key"]
            for k in await ops.chat_epoch_keys(node["state"], GROUP)}
    assert len(keys) == 2
    texts = []
    for m in await node["chat"].get_recent(10):
        texts.append(open_message(keys[m.epoch], GROUP, m.epoch, b64,
                                  m.nonce, m.payload)["text"])
    assert texts == ["epoch one", "epoch two"]


async def test_an_epoch_key_is_never_written_in_the_clear(node):
    """
    R15. The claim chat encryption makes is against someone who takes the
    node's storage *without the keystore password*. An epoch key sitting in a
    plaintext SQLite beside chat.db would collapse that to nothing, silently,
    and it is the obvious thing to write.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    keys = await ops.chat_epoch_keys(node["state"], GROUP)
    assert keys

    live = keys[-1]["key"]
    for path in sorted(node["tmp_path"].rglob("*")):
        if not path.is_file():
            continue
        assert live not in path.read_bytes(), (
            f"the live chat epoch key appears verbatim in {path.name} — "
            "it must be wrapped to the node's own key, as the GEK is")


async def test_the_stored_message_contains_neither_text_nor_display_name(node):
    """
    What "encrypted at rest" has to mean. The display name is inside the
    envelope too: on the wire it is a field any peer can set to anything, and
    the node caches it to render history, so leaving it outside would both
    leak it and leave spoofing free.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    sk, raw, b64 = _device()
    session = _session(node, device_b64=b64)
    await _send_sealed(node, session, sk, raw, b64, "a secret message")

    blob = (node["tmp_path"] / "chat.db").read_bytes()
    assert b"a secret message" not in blob
    stored = (await node["chat"].get_recent(10))[0]
    assert stored.format == FORMAT_SEALED_V1
    assert b"a secret message" not in stored.payload


# ── refusals ─────────────────────────────────────────────────────────────────

async def test_plaintext_is_refused_always(node):
    """
    R5 / C6's lesson one feature later, and now unconditional: there is no
    switch to leave in the wrong position. A member who can post in clear into
    a group whose members believe their chat is encrypted is a downgrade anyone
    could ask for.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    session = _session(node)
    _spawn_inline(session)
    session._do_chat_message({"payload": "in the clear", "sender_name": "alice"})

    assert session.sent[-1]["type"] == "error"
    assert await node["chat"].message_count() == 0


async def test_there_is_no_setting_that_re_enables_plaintext(node):
    """
    The switch is gone, not defaulted. A `chat_encrypted` in the group context
    — left by an older node's roster row, or invented by anything reading one —
    must not be consulted, or the bypass is back with a name.
    """
    node["group_ctx"]["chat_encrypted"] = False
    await ops.ensure_chat_epoch(node["state"], GROUP)
    session = _session(node)
    _spawn_inline(session)
    session._do_chat_message({"payload": "in the clear", "sender_name": "alice"})

    assert session.sent[-1]["type"] == "error"
    assert await node["chat"].message_count() == 0

    source = (Path(__file__).parent.parent / "src" / "meshbay_node"
              / "transport" / "webrtc_server.py").read_text(encoding="utf-8")
    assert 'get("chat_encrypted"' not in source, (
        "nothing may read a chat_encrypted setting — there is no switch")


async def test_a_member_cannot_send_as_another_members_device(node):
    """
    The hole that would have made encrypted chat *worse* than plaintext chat.

    Receivers verify a signature against the `device` field, so a member free
    to name somebody else's key could be that member to everyone — and the
    signature would check out. The connection has proved which device it is,
    and the claim must match it.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    _sk_alice, raw_alice, b64_alice = _device()
    sk_mallory, raw_mallory, b64_mallory = _device()

    session = _session(node, user_id="mallory", device_b64=b64_mallory)
    keys = await ops.chat_epoch_keys(node["state"], GROUP)
    epoch, key = keys[-1]["epoch"], keys[-1]["key"]
    # Correctly sealed and correctly signed — by Mallory, claiming to be Alice.
    env = seal(key, GROUP, epoch, b64_alice, raw_alice, sk_mallory,
               {"text": "not from alice"})
    _spawn_inline(session)
    session._do_chat_message({
        "format": FORMAT_SEALED_V1, "epoch": epoch, "device": raw_alice,
        "ct": env["ct"], "nonce": env["nonce"], "sig": env["sig"],
    })

    assert session.sent[-1]["type"] == "error"
    assert await node["chat"].message_count() == 0


async def test_a_signed_message_cannot_be_replayed(node):
    """
    A replay is a *validly signed* copy of a real message, so nothing about
    the signature refuses it. The unique (device, nonce) does — and the nonce
    is already required to be unique for AES-GCM to be safe, so it costs
    nothing to make it a key.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    sk, raw, b64 = _device()
    session = _session(node, device_b64=b64)
    env = await _send_sealed(node, session, sk, raw, b64, "said once")
    assert await node["chat"].message_count() == 1

    keys = await ops.chat_epoch_keys(node["state"], GROUP)
    pending = _spawn_inline(session)
    session._do_chat_message({
        "format": FORMAT_SEALED_V1, "epoch": keys[-1]["epoch"], "device": raw,
        "ct": env["ct"], "nonce": env["nonce"], "sig": env["sig"],
    })
    for task in pending:
        await task
    assert await node["chat"].message_count() == 1, (
        "a replayed message must not be stored twice")


async def test_an_unidentified_connection_cannot_send_a_signed_message(node):
    """
    `device_hello` is what makes "that is not the device on this connection"
    checkable at all. Without it the node knows the account and not the key,
    and a `device` field would be an assertion nobody verified.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    sk, raw, b64 = _device()
    session = _session(node)              # no device_hello
    keys = await ops.chat_epoch_keys(node["state"], GROUP)
    epoch, key = keys[-1]["epoch"], keys[-1]["key"]
    env = seal(key, GROUP, epoch, b64, raw, sk, {"text": "x"})
    _spawn_inline(session)
    session._do_chat_message({
        "format": FORMAT_SEALED_V1, "epoch": epoch, "device": raw,
        "ct": env["ct"], "nonce": env["nonce"], "sig": env["sig"],
    })
    assert session.sent[-1]["type"] == "error"


# ── key delivery ─────────────────────────────────────────────────────────────

async def test_the_keys_are_delivered_sealed_under_the_group_key(node):
    """
    Sealed for the same reason the index and the ack are, one step stronger:
    the payload *is* key material. A peer that has completed the handshake
    holds the group key and can open it; anything short of that gets a
    ciphertext.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    await ops.open_chat_epoch(node["state"], GROUP)
    session = _session(node)
    await session._do_chat_keys_req({})

    resp = session.sent[-1]
    assert resp["type"] == MNP.CHAT_KEYS_RESP
    assert "epochs" not in resp, "the keys must not travel in clear"
    payload = unseal(node["gek"], PURPOSE_CHAT_KEYS, MNP.CHAT_KEYS_RESP,
                     GROUP, resp)
    assert [e["epoch"] for e in payload["epochs"]] == [1, 2]
    assert payload["current"] == 2
    for e in payload["epochs"]:
        assert len(e["key"]) == 32


async def test_every_epoch_is_delivered_not_just_the_current_one(node):
    """
    R2 again, from the delivery side: this is what lets a device linked this
    morning read a conversation from last year.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    for _ in range(3):
        await ops.open_chat_epoch(node["state"], GROUP)
    session = _session(node)
    await session._do_chat_keys_req({})
    payload = unseal(node["gek"], PURPOSE_CHAT_KEYS, MNP.CHAT_KEYS_RESP,
                     GROUP, session.sent[-1])
    assert [e["epoch"] for e in payload["epochs"]] == [1, 2, 3, 4]


# ── epochs move when access shrinks ─────────────────────────────────────────

async def test_revoking_a_device_opens_a_new_epoch(node):
    """
    A revoked device holds every chat key it ever received. Revocation stops
    the node handing over the *next* one; nothing else takes the current one
    away — the exact counterpart of "still rotate the GEK".
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    before = await node["bundles"].latest_chat_epoch(GROUP)
    session = _session(node)
    await session._new_chat_epoch(GROUP, "device_revoke")
    assert await node["bundles"].latest_chat_epoch(GROUP) == before + 1


async def test_a_group_always_gets_an_epoch(node):
    """
    Chat is always encrypted, so a group with no epoch key is a group nobody
    can speak in. `ensure_chat_epoch` is what the daemon calls at group load —
    at start-up, where a failure lands in the log the operator is already
    reading rather than on somebody's first message.
    """
    assert await node["bundles"].latest_chat_epoch(GROUP) == 0
    epoch = await ops.ensure_chat_epoch(node["state"], GROUP)
    assert epoch == 1
    # Idempotent: called at every group load, and a second epoch per restart
    # would be a key nobody needed and the node keeps for ever.
    assert await ops.ensure_chat_epoch(node["state"], GROUP) == 1


async def test_an_epoch_key_is_never_deleted(node):
    """
    Nothing in the system removes an epoch key, and nothing may: the messages
    sealed under it become unreadable the moment it goes, for everybody. The
    only operation that touches the table adds a row.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    sk, raw, b64 = _device()
    session = _session(node, device_b64=b64)
    await _send_sealed(node, session, sk, raw, b64, "still readable")
    await ops.open_chat_epoch(node["state"], GROUP)
    await ops.prune_chat(node["state"], GROUP, 3650)

    keys = await ops.chat_epoch_keys(node["state"], GROUP)
    assert [k["epoch"] for k in keys] == [1, 2]

    source = (Path(__file__).parent.parent / "src" / "meshbay_node"
              / "bundle_store.py").read_text(encoding="utf-8")
    assert "DELETE FROM chat_epochs" not in source
    assert "INSERT OR REPLACE INTO chat_epochs" not in source, (
        "an epoch key is written once — REPLACE would destroy the history "
        "sealed under it, with no error anywhere")


# ── the explicit history migration, and retention ───────────────────────────

async def test_encrypt_history_converts_the_old_plaintext(node):
    """
    The migration for a node that ran before MNP 2.0.

    The plaintext row is written straight into the store, because that is the
    only way one can exist now: `_do_chat_message` refuses plaintext outright.
    Such rows are the ones still readable off a stolen disk, and the node can
    convert them only because it holds them in the clear — it is the last
    moment at which anyone can.
    """
    node["state"]["sk_node"] = Ed25519PrivateKey.generate()
    await node["chat"].save_message(
        sender_id="alice", iteration=0, payload=b"written in the clear",
        sender_name="alice")
    await ops.ensure_chat_epoch(node["state"], GROUP)

    result = await ops.encrypt_chat_history(node["state"], GROUP)

    assert result["converted"] == 1
    stored = (await node["chat"].get_recent(10))[0]
    assert stored.format == FORMAT_SEALED_V1
    assert b"written in the clear" not in stored.payload
    assert stored.sender_name == "", (
        "the display name moves inside the envelope — leaving it would keep in "
        "the clear the one field the sealing was for")

    keys = {k["epoch"]: k["key"]
            for k in await ops.chat_epoch_keys(node["state"], GROUP)}
    device_b64 = base64.b64encode(stored.device).decode()
    opened = open_message(keys[stored.epoch], GROUP, stored.epoch, device_b64,
                          stored.nonce, stored.payload)
    assert opened["text"] == "written in the clear"
    assert opened["sender_name"] == "alice"
    assert opened["migrated"] is True, (
        "a migrated message carries the node's word for who wrote it, which is "
        "all it ever carried — that has to be visible, not inferred")


async def test_encrypt_history_backs_the_database_up_first(node):
    node["state"]["sk_node"] = Ed25519PrivateKey.generate()
    await node["chat"].save_message(
        sender_id="alice", iteration=0, payload=b"one", sender_name="alice")
    await ops.ensure_chat_epoch(node["state"], GROUP)

    result = await ops.encrypt_chat_history(node["state"], GROUP)

    from pathlib import Path
    backup = Path(result["backup"])
    assert backup.exists() and backup.stat().st_size > 0
    assert b"one" in backup.read_bytes(), (
        "the backup is taken before the rewrite, or it is not a backup")


async def test_retention_deletes_messages_and_never_epoch_keys(node):
    """
    R16. An epoch whose messages have all aged out costs 32 bytes; deleting it
    would make anything still stored under it unreadable.
    """
    await ops.ensure_chat_epoch(node["state"], GROUP)
    sk, raw, b64 = _device()
    session = _session(node, device_b64=b64)
    await _send_sealed(node, session, sk, raw, b64, "old news")

    # Age it past the cutoff.
    await node["chat"]._db.execute(
        "UPDATE messages SET timestamp = ?", (time.time() - 40 * 86400,))
    await node["chat"].commit()

    result = await ops.prune_chat(node["state"], GROUP, 30)

    assert result["removed"] == 1
    assert await node["chat"].message_count() == 0
    assert await ops.chat_epoch_keys(node["state"], GROUP), (
        "retention deletes messages, never keys")


async def test_retention_refuses_a_zero_day_window(node):
    """`prune 0` would delete the whole conversation and read as a typo."""
    with pytest.raises(ops.OpError):
        await ops.prune_chat(node["state"], GROUP, 0)