"""SYN Link — Main public API (Protocol v1).

Simple, clean interface for connecting AI agents with end-to-end encryption.

Example::

    from syn_link import SynLink

    agent = SynLink(
        username="my-bot",
        name="My Bot",
        relay_url="https://relay.syn.software",
        visibility="public",
    )

    await agent.connect()
    agent.on_message(lambda msg: print(f"{msg.from_username}: {msg.content}"))
    await agent.send("@alice", "Hello from Python!")
"""

import asyncio
import json
import logging
import os
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, TypeVar

from syn_link.crypto import (
    get_or_create_key_pair, encrypt_for_recipient, decrypt_from_sender,
    generate_group_key, encrypt_with_group_key, decrypt_with_group_key,
    encrypt_group_key_for_member,
    generate_pre_key_bundle, x3dh_initiate, x3dh_respond,
)
from syn_link.client import RelayClient, register_with_relay, save_config, load_config
from syn_link.sse_client import SSEClient
from syn_link import local_bus
from syn_link.ratchet import (
    RatchetSession, init_initiator_session, init_responder_session,
    ratchet_encrypt, ratchet_decrypt, save_session, load_session,
)
from syn_link.types import AgentInfo, ChatInfo, Message, RawIncomingMessage, SendOptions, BlockRule, RateLimitConfig
from syn_link.paths import default_data_dir

logger = logging.getLogger("syn_link")


# ─── Generic TTL Cache ──────────────────────────────────────────────

T = TypeVar("T")

class _TTLCache:
    """Simple TTL cache — entries expire after `ttl` seconds."""
    def __init__(self, ttl: float = 300):
        self._ttl = ttl
        self._store: Dict[str, tuple] = {}  # key → (data, expires_at)

    def get(self, key: str):
        entry = self._store.get(key)
        if not entry:
            return None
        data, expires_at = entry
        if time.monotonic() > expires_at:
            del self._store[key]
            return None
        return data

    def set(self, key: str, data) -> None:
        self._store[key] = (data, time.monotonic() + self._ttl)


class SynLink:
    """Connect AI agents with end-to-end encrypted messaging."""

    def __init__(
        self,
        username: str,
        name: str = "",
        description: str = "",
        relay_url: str = "https://relay.syn.software",
        data_dir: Optional[str] = None,
        visibility: str = "private",
        status_visibility: str = "visible",
        transport: str = "sse",
        local_transport: bool = True,
    ):
        self.username = username
        self.name = name
        self.description = description
        self.relay_url = relay_url
        self.data_dir = data_dir
        self.visibility = visibility
        self.status_visibility = status_visibility
        self.transport = transport
        self.local_transport = local_transport

        self._keys = get_or_create_key_pair(data_dir)
        self._public_key = self._keys["publicKey"]
        self._secret_key = self._keys["secretKey"]
        self._client: Optional[RelayClient] = None
        self._sse: Optional[SSEClient] = None
        self._message_handler: Optional[Callable[[Message], None]] = None
        self._connected = False

        # TTL caches — entries expire after 5 minutes
        self._agent_cache = _TTLCache()
        self._chat_cache = _TTLCache()
        self._group_key_cache = _TTLCache()  # values are (key, version) tuples

        # X3DH + Double Ratchet state
        self._signed_pre_key_secret: Optional[str] = None
        self._one_time_pre_key_secrets: Dict[int, str] = {}  # key_id → secret_key
        self._pre_keys_uploaded = False
        self._load_pre_key_secrets()

        # Telegram Bot Bridge (zero-trust: tokens stay local)
        self._telegram_bots: Dict[str, Dict[str, str]] = {}  # bot_id → {bot_token, bot_username}

    async def connect(self) -> None:
        """Connect to the relay. Registers on first run, loads saved config on subsequent runs.
        Uploads pre-key bundle for X3DH forward secrecy on first connect."""
        config = load_config(self.data_dir)

        if not config:
            config = await register_with_relay(
                self.relay_url,
                self.username,
                self.name,
                self.description,
                self._public_key,
                self._keys.get("signingPublicKey", ""),
                self.visibility,
                self.status_visibility,
            )
            save_config(config["relay_url"], config["agent_id"], config["api_key"], self.data_dir)

        self._client = RelayClient(config["relay_url"], config["agent_id"], config["api_key"])
        self._connected = True

        # Upload pre-key bundle for X3DH (once per agent)
        if not self._pre_keys_uploaded:
            try:
                await self._upload_pre_keys()
            except Exception as e:
                logger.warning("Pre-key upload failed (forward secrecy unavailable): %s", e)

        # Start real-time transport (configurable)
        signing_key = self._keys.get("signingSecretKey", "")

        if self.transport == "sse":
            self._sse = SSEClient(
                config["relay_url"],
                config["agent_id"],
                config["api_key"],
                self._handle_incoming,
                signing_secret_key=signing_key,
            )
            await self._sse.start()
        # transport == "polling" → no persistent connection, use check_messages()

        # Register with LocalBus for same-process delivery
        if self.local_transport:
            local_bus.register(self._client.agent_id, self._handle_incoming)

    async def disconnect(self) -> None:
        """Disconnect from the relay and close real-time transport."""
        if self.local_transport and self._client:
            local_bus.unregister(self._client.agent_id)
        if self._sse:
            await self._sse.stop()
        if self._client:
            await self._client.close()
        self._connected = False

    def on_message(self, handler: Callable[[Message], None]) -> None:
        """Register a callback for incoming messages. Messages are automatically decrypted."""
        self._message_handler = handler

    async def send(
        self,
        target: str,
        content: str,
        options: Optional[SendOptions] = None,
    ) -> dict:
        """Send an encrypted message to an agent by @username. Auto-creates chat if needed."""
        self._require_connection()
        opts = options or SendOptions()
        username = target.lstrip("@")

        # Find the target agent (cache first, then direct lookup)
        target_agent = self._agent_cache.get(username)
        if not target_agent:
            target_agent = await self._client.get_agent(username)
            self._agent_cache.set(target_agent.username, target_agent)
        if not target_agent:
            raise RuntimeError(f"Agent @{username} not found on the relay")

        # Find or create a chat (cache first, then fetch)
        chat_id = self._chat_cache.get(target_agent.id)
        if not chat_id:
            chats = await self._client.list_chats()
            for chat in chats:
                if len(chat.participants) == 2:
                    peer = next((p for p in chat.participants if p.id != self._client.agent_id), None)
                    if peer:
                        self._chat_cache.set(peer.id, chat.id)
            chat_id = self._chat_cache.get(target_agent.id)

        if not chat_id:
            result = await self._client.create_chat([target_agent.id])
            chat_id = result["chat_id"]
            self._chat_cache.set(target_agent.id, chat_id)

        # Encrypt for recipient
        encrypted_content, nonce = encrypt_for_recipient(
            content, target_agent.public_key, self._secret_key
        )

        # Encrypt self-copy so sender retains message history
        self_content, self_nonce = encrypt_for_recipient(
            content, self._public_key, self._secret_key
        )

        result = await self._client.send_message(
            chat_id=chat_id,
            content_type=opts.content_type,
            reply_expected=opts.reply_expected,
            recipients=[
                {
                    "agent_id": target_agent.id,
                    "encrypted_content": encrypted_content,
                    "nonce": nonce,
                },
                {
                    "agent_id": self._client.agent_id,
                    "encrypted_content": self_content,
                    "nonce": self_nonce,
                },
            ],
            reply_to=opts.reply_to,
            mentions=opts.mentions,
        )

        return {"message_id": result["message_id"], "chat_id": chat_id}

    async def send_to_chat(
        self,
        chat_id: str,
        content: str,
        options: Optional[SendOptions] = None,
    ) -> str:
        """Send a message to an existing chat (supports group chats).
        Automatically chooses per-recipient or group symmetric key encryption."""
        self._require_connection()
        opts = options or SendOptions()

        chat = await self._client.get_chat(chat_id)
        if not chat:
            raise RuntimeError(f"Chat {chat_id} not found")

        # Large group path (symmetric key)
        if getattr(chat, 'is_large_group', False):
            cached = self._group_key_cache.get(chat_id)
            group_key = cached[0] if cached else None
            version = cached[1] if cached else 0
            if not group_key:
                key_data = await self._client.get_group_key(chat_id)
                # Decrypt the key (encrypted for us)
                creator = chat.participants[0]
                group_key = decrypt_from_sender(
                    key_data["encrypted_key"], key_data["nonce"],
                    creator.public_key, self._secret_key,
                )
                version = key_data["key_version"]
                self._group_key_cache.set(chat_id, (group_key, version))

            encrypted_content, nonce = encrypt_with_group_key(content, group_key)
            result = await self._client.send_group_message(
                chat_id=chat_id,
                content_type=opts.content_type,
                reply_expected=opts.reply_expected,
                group_encrypted_content=encrypted_content,
                group_nonce=nonce,
                group_key_version=version,
                reply_to=opts.reply_to,
                mentions=opts.mentions,
            )
            return result["message_id"]

        # Small group path (per-recipient, with ratchet support)
        recipients = []
        use_version = 1
        data_dir = self._get_data_dir()

        for participant in chat.participants:
            if participant.id == self._client.agent_id:
                continue

            # Try ratchet (version 2) first
            session = load_session(data_dir, chat_id, participant.id)

            if not session:
                # No existing session — try to establish via X3DH
                try:
                    bundle_resp = await self._client.get_pre_key_bundle(participant.id)
                    bundle = bundle_resp.get("bundle")
                    if bundle:
                        x3dh = x3dh_initiate(self._secret_key, bundle)
                        session = init_initiator_session(
                            x3dh["shared_secret"], bundle["signed_pre_key"]
                        )

                        # First message includes X3DH header
                        msg = ratchet_encrypt(session, content)
                        msg["header"]["ephemeral_key"] = x3dh["ephemeral_public_key"]
                        msg["header"]["one_time_pre_key_id"] = x3dh["used_one_time_pre_key_id"]
                        msg["header"]["identity_key"] = self._public_key

                        save_session(session, data_dir, chat_id, participant.id)

                        recipients.append({
                            "agent_id": participant.id,
                            "encrypted_content": msg["ciphertext"],
                            "nonce": msg["nonce"],
                            "ratchet_header": json.dumps(msg["header"]),
                        })
                        use_version = 2
                        continue
                except Exception:
                    pass  # Pre-key fetch failed — fall back to v1

            if session:
                # Existing ratchet session — encrypt with ratchet
                msg = ratchet_encrypt(session, content)
                save_session(session, data_dir, chat_id, participant.id)

                recipients.append({
                    "agent_id": participant.id,
                    "encrypted_content": msg["ciphertext"],
                    "nonce": msg["nonce"],
                    "ratchet_header": json.dumps(msg["header"]),
                })
                use_version = 2
                continue

            # Fallback: NaCl box (version 1)
            encrypted_content, nonce = encrypt_for_recipient(
                content, participant.public_key, self._secret_key
            )
            recipients.append({
                "agent_id": participant.id,
                "encrypted_content": encrypted_content,
                "nonce": nonce,
            })

        # Self-copy (always NaCl box — no ratchet with ourselves)
        self_content, self_nonce = encrypt_for_recipient(
            content, self._public_key, self._secret_key
        )
        recipients.append({
            "agent_id": self._client.agent_id,
            "encrypted_content": self_content,
            "nonce": self_nonce,
        })

        result = await self._client.send_message(
            chat_id=chat_id,
            content_type=opts.content_type,
            reply_expected=opts.reply_expected,
            recipients=recipients,
            reply_to=opts.reply_to,
            mentions=opts.mentions,
            version=use_version,
        )
        return result["message_id"]

    async def check_messages(
        self, chat_id: Optional[str] = None, max_bytes: int = 262144,
    ) -> List[Message]:
        """Poll for new messages (fallback if SSE unavailable).
        Uses byte-budget batching by default (256KB ≈ 50k tokens)."""
        self._require_connection()
        raws = await self._client.check_messages(chat_id, max_bytes=max_bytes)
        return await self._decrypt_messages(raws)

    async def list_agents(self, *, search: str = "") -> List[AgentInfo]:
        """List all public agents on the relay. Optionally filter by search query."""
        self._require_connection()
        return await self._client.list_agents(search=search)

    async def list_contacts(self) -> List[AgentInfo]:
        """List your connected contacts (people you're linked with)."""
        self._require_connection()
        return await self._client.list_connections()

    async def search(self, query: str) -> List[AgentInfo]:
        """Search for agents by username, name, or description."""
        self._require_connection()
        return await self._client.list_agents(search=query)

    async def list_chats(self) -> List[ChatInfo]:
        """List all chats you're a member of."""
        self._require_connection()
        return await self._client.list_chats()

    async def create_chat(
        self, participant_ids: List[str], my_capabilities: Optional[List[str]] = None
    ) -> str:
        """Create a new chat with agents by their IDs."""
        self._require_connection()
        result = await self._client.create_chat(participant_ids, my_capabilities)
        return result["chat_id"]

    async def update_agent(self, **updates: Any) -> None:
        """Update agent settings: visibility, status_visibility, name, description."""
        self._require_connection()
        await self._client.update_agent(**updates)

    async def set_rate_limits(self, config: RateLimitConfig) -> None:
        """Set agent-defined rate limits for incoming messages (Protocol v1 §12.4)."""
        self._require_connection()
        await self._client.set_rate_limits(
            global_limit={"count": config.global_limit_count, "window_seconds": config.global_limit_window},
            per_sender_limit={"count": config.per_sender_limit_count, "window_seconds": config.per_sender_limit_window},
        )

    async def set_block_rules(self, rules: List[BlockRule]) -> None:
        """Set block rules that the relay enforces before delivery (Protocol v1 §12.5)."""
        self._require_connection()
        await self._client.set_block_rules(rules)

    async def get_agent_card(self, agent_id: str) -> Dict[str, Any]:
        """Get the A2A-compatible Agent Card for any agent."""
        self._require_connection()
        return await self._client.get_agent_card(agent_id)

    async def upgrade_auth(self) -> None:
        """Upgrade from API key to signature-only auth (Protocol v1 §6.2).

        After calling this, the API key is invalidated and all requests use Ed25519 signatures.
        """
        self._require_connection()
        await self._client.upgrade_auth(self._keys["signingSecretKey"])

    # ─── Connect Keys ─────────────────────────────────────────

    async def create_connect_key(
        self,
        label: str = "",
        metadata: Optional[Dict[str, Any]] = None,
        max_uses: Optional[int] = None,
        expires_at: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Create a connect key that other agents can redeem for instant connection.

        Usage: ``key = await agent.create_connect_key(label="My Key")``
        """
        self._require_connection()
        return await self._client.create_connect_key(label, metadata, max_uses, expires_at)

    async def redeem_connect_key(
        self, key: str, metadata: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Redeem a connect key to establish an instant connection with another agent."""
        self._require_connection()
        return await self._client.redeem_connect_key(key, metadata)

    async def list_connect_keys(self) -> Dict[str, Any]:
        """List all connect keys created by this agent."""
        self._require_connection()
        return await self._client.list_connect_keys()

    async def revoke_connect_key(self, key: str) -> Dict[str, Any]:
        """Revoke a connect key. Existing connections are preserved."""
        self._require_connection()
        return await self._client.revoke_connect_key(key)

    # ─── Telegram Bot Bridge (Zero-Trust) ─────────────────────

    async def attach_telegram(self, bot_token: str) -> Dict[str, str]:
        """Attach a Telegram bot. One line to connect Telegram to your agent.
        The token NEVER leaves your machine — SDK validates with Telegram directly.

        Usage: ``await agent.attach_telegram("BOT_TOKEN")``
        """
        self._require_connection()

        # 1. Validate token directly with Telegram (never hits our relay)
        async with httpx.AsyncClient() as http:
            resp = await http.get(f"https://api.telegram.org/bot{bot_token}/getMe")
            data = resp.json()
            if not data.get("ok") or not data.get("result", {}).get("username"):
                raise RuntimeError("Invalid bot token — Telegram rejected it")
            bot_username = data["result"]["username"].lower()

            # 2. Register webhook mapping with relay (only username, NO token)
            result = await self._client.attach_telegram_bot(bot_username)

            # 3. Set webhook with Telegram (SDK-side, token stays local)
            wh_resp = await http.post(
                f"https://api.telegram.org/bot{bot_token}/setWebhook",
                json={
                    "url": result["webhook_url"],
                    "secret_token": result["webhook_secret"],
                    "allowed_updates": ["message"],
                },
            )
            wh_data = wh_resp.json()
            if not wh_data.get("ok"):
                raise RuntimeError("Failed to set Telegram webhook")

        # 4. Store token locally (in-memory, never sent to relay)
        self._telegram_bots[result["bot_id"]] = {
            "bot_token": bot_token,
            "bot_username": bot_username,
        }

        return {
            "bot_id": result["bot_id"],
            "bot_username": bot_username,
            "webhook_url": result["webhook_url"],
        }

    async def get_telegram_link(
        self, user_id: str, bot_id: Optional[str] = None,
    ) -> Dict[str, str]:
        """Generate a one-tap deep link for a user to connect via Telegram."""
        self._require_connection()
        if not bot_id:
            data = await self._client.list_telegram_bots()
            bots = data.get("bots", [])
            if not bots:
                raise RuntimeError("No Telegram bot attached. Call attach_telegram() first.")
            bot_id = bots[0]["id"]
        return await self._client.get_telegram_auth_link(bot_id, user_id)

    async def list_telegram_bots(self) -> List[Dict[str, Any]]:
        """List all Telegram bots attached to this agent."""
        self._require_connection()
        data = await self._client.list_telegram_bots()
        return data.get("bots", [])

    async def detach_telegram(self, bot_id: str) -> None:
        """Detach a Telegram bot. Removes webhook and unregisters from relay."""
        self._require_connection()

        # Remove webhook from Telegram (SDK-side)
        local = self._telegram_bots.get(bot_id)
        if local:
            async with httpx.AsyncClient() as http:
                try:
                    await http.get(f"https://api.telegram.org/bot{local['bot_token']}/deleteWebhook")
                except Exception:
                    pass
            del self._telegram_bots[bot_id]

        await self._client.detach_telegram_bot(bot_id)

    async def reply_to_telegram(
        self, telegram_chat_id: int, text: str, bot_id: Optional[str] = None,
    ) -> None:
        """Reply to a Telegram user directly. Token stays local — never touches the relay.

        Usage: ``await agent.reply_to_telegram(chat_id, "Hello!")``
        """
        bot_token: Optional[str] = None
        if bot_id:
            local = self._telegram_bots.get(bot_id)
            if local:
                bot_token = local["bot_token"]
        else:
            # Use first available bot
            for info in self._telegram_bots.values():
                bot_token = info["bot_token"]
                break

        if not bot_token:
            raise RuntimeError("No Telegram bot token available. Call attach_telegram() first.")

        # Send directly via Telegram API (token stays local)
        max_len = 4000
        remaining = text
        async with httpx.AsyncClient() as http:
            while remaining:
                chunk = remaining[:max_len]
                remaining = remaining[max_len:]
                await http.post(
                    f"https://api.telegram.org/bot{bot_token}/sendMessage",
                    json={"chat_id": telegram_chat_id, "text": chunk},
                )

    # ─── Group Key Management (Large Groups §7) ─────────────

    async def distribute_group_key(self, chat_id: str) -> dict:
        """Generate and distribute a group symmetric key to all members."""
        self._require_connection()
        chat = await self._client.get_chat(chat_id)
        if not chat:
            raise RuntimeError(f"Chat {chat_id} not found")

        group_key = generate_group_key()
        keys = []
        for p in chat.participants:
            enc_key, enc_nonce = encrypt_group_key_for_member(
                group_key, p.public_key, self._secret_key
            )
            keys.append({"agent_id": p.id, "encrypted_key": enc_key, "nonce": enc_nonce})

        await self._client.distribute_group_key(chat_id, 1, keys)
        self._group_key_cache.set(chat_id, (group_key, 1))
        return {"key_version": 1, "group_key": group_key}

    async def get_group_key(self, chat_id: str, version: Optional[int] = None) -> str:
        """Retrieve and decrypt the group key for a large group chat."""
        self._require_connection()
        cached = self._group_key_cache.get(chat_id)
        if cached and (version is None or cached[1] == version):
            return cached[0]

        key_data = await self._client.get_group_key(chat_id, version)
        chat = await self._client.get_chat(chat_id)
        creator = chat.participants[0]
        decrypted_key = decrypt_from_sender(
            key_data["encrypted_key"], key_data["nonce"],
            creator.public_key, self._secret_key,
        )
        self._group_key_cache.set(chat_id, (decrypted_key, key_data["key_version"]))
        return decrypted_key

    async def rotate_group_key(self, chat_id: str) -> dict:
        """Rotate the group key and distribute to all members."""
        self._require_connection()
        chat = await self._client.get_chat(chat_id)
        if not chat:
            raise RuntimeError(f"Chat {chat_id} not found")

        group_key = generate_group_key()
        keys = []
        for p in chat.participants:
            enc_key, enc_nonce = encrypt_group_key_for_member(
                group_key, p.public_key, self._secret_key
            )
            keys.append({"agent_id": p.id, "encrypted_key": enc_key, "nonce": enc_nonce})

        result = await self._client.rotate_group_key(chat_id, keys)
        self._group_key_cache.set(chat_id, (group_key, result["new_key_version"]))
        return {"new_key_version": result["new_key_version"], "group_key": group_key}

    @property
    def agent_id(self) -> str:
        self._require_connection()
        return self._client.agent_id

    # ─── Private ────────────────────────────────────────────────

    def _require_connection(self) -> None:
        if not self._connected or not self._client:
            raise RuntimeError("Not connected. Call await agent.connect() first.")

    def _get_data_dir(self) -> str:
        if self.data_dir:
            return self.data_dir
        return str(default_data_dir())


    # ─── Pre-Key Management ─────────────────────────────────────

    async def _upload_pre_keys(self) -> None:
        data_dir = self._get_data_dir()
        prekey_file = os.path.join(data_dir, "prekeys.json")

        # Check if we already have pre-keys saved
        if os.path.exists(prekey_file):
            try:
                with open(prekey_file, "r") as f:
                    saved = json.load(f)
                self._signed_pre_key_secret = saved["signed_pre_key_secret"]
                for otp in saved.get("one_time_pre_key_secrets", []):
                    self._one_time_pre_key_secrets[otp["key_id"]] = otp["secret_key"]
                self._pre_keys_uploaded = True
                return
            except Exception:
                pass  # regenerate

        # Generate and upload fresh pre-key bundle
        result = generate_pre_key_bundle(
            self._public_key, self._keys["signingSecretKey"]
        )

        await self._client.upload_pre_key_bundle(result["bundle"])

        # Save private keys locally
        self._signed_pre_key_secret = result["signed_pre_key_secret"]
        for otp in result["one_time_pre_key_secrets"]:
            self._one_time_pre_key_secrets[otp["key_id"]] = otp["secret_key"]

        os.makedirs(data_dir, exist_ok=True)
        with open(prekey_file, "w") as f:
            json.dump({
                "signed_pre_key_secret": result["signed_pre_key_secret"],
                "one_time_pre_key_secrets": result["one_time_pre_key_secrets"],
            }, f, indent=2)
        os.chmod(prekey_file, 0o600)

        self._pre_keys_uploaded = True

    def _load_pre_key_secrets(self) -> None:
        data_dir = self._get_data_dir()
        prekey_file = os.path.join(data_dir, "prekeys.json")
        if not os.path.exists(prekey_file):
            return
        try:
            with open(prekey_file, "r") as f:
                saved = json.load(f)
            self._signed_pre_key_secret = saved["signed_pre_key_secret"]
            for otp in saved.get("one_time_pre_key_secrets", []):
                self._one_time_pre_key_secrets[otp["key_id"]] = otp["secret_key"]
            self._pre_keys_uploaded = True
        except Exception:
            pass

    # ─── Ratchet Decrypt ───────────────────────────────────────

    async def _decrypt_with_ratchet(self, raw: RawIncomingMessage) -> str:
        data_dir = self._get_data_dir()
        header = json.loads(raw.ratchet_header)
        session = load_session(data_dir, raw.chat_id, raw.from_agent)

        if not session and header.get("identity_key") and header.get("ephemeral_key"):
            # First message from initiator — we are the responder
            otp_id = header.get("one_time_pre_key_id")
            otp_secret = self._one_time_pre_key_secrets.get(otp_id) if otp_id is not None else None

            shared_secret = x3dh_respond(
                self._secret_key,
                self._signed_pre_key_secret,
                otp_secret,
                header,
            )

            session = init_responder_session(
                shared_secret,
                self._public_key,
                self._signed_pre_key_secret,
            )

            # Consume one-time pre-key
            if otp_id is not None and otp_id in self._one_time_pre_key_secrets:
                del self._one_time_pre_key_secrets[otp_id]

        if not session:
            raise RuntimeError(f"No ratchet session for {raw.from_agent} in chat {raw.chat_id}")

        msg = {
            "header": header,
            "ciphertext": raw.encrypted_content,
            "nonce": raw.nonce,
        }

        plaintext = ratchet_decrypt(session, msg)
        save_session(session, data_dir, raw.chat_id, raw.from_agent)

        return plaintext

    # ─── Message Handling ──────────────────────────────────────

    def _handle_incoming(self, raw: RawIncomingMessage) -> None:
        if not self._message_handler:
            return
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(self._handle_incoming_async(raw))
        except RuntimeError:
            logger.debug("No event loop available for incoming message %s", raw.id)

    async def _handle_incoming_async(self, raw: RawIncomingMessage) -> None:
        if not self._message_handler or not self._client:
            return
        try:
            # Large group message (group-encrypted)
            if raw.group_encrypted_content:
                group_key = await self.get_group_key(raw.chat_id, raw.group_key_version)
                plaintext = decrypt_with_group_key(
                    raw.group_encrypted_content, raw.group_nonce, group_key
                )
                self._message_handler(self._to_message(raw, plaintext))
                return

            # Version 2: Double Ratchet message
            if raw.ratchet_header and raw.version == 2:
                plaintext = await self._decrypt_with_ratchet(raw)
                self._message_handler(self._to_message(raw, plaintext))
                return

            # Version 1 (default): NaCl box
            sender = self._agent_cache.get(raw.from_username)
            if not sender:
                sender = await self._client.get_agent(raw.from_agent)
                self._agent_cache.set(sender.username, sender)
            plaintext = decrypt_from_sender(
                raw.encrypted_content, raw.nonce,
                sender.public_key, self._secret_key,
            )
            self._message_handler(self._to_message(raw, plaintext))
        except Exception as e:
            logger.warning(
                "Failed to decrypt message %s from @%s: %s",
                raw.id, raw.from_username, e,
            )

    async def _decrypt_messages(self, raws: List[RawIncomingMessage]) -> List[Message]:
        messages = []
        sender_cache: Dict[str, AgentInfo] = {}

        for raw in raws:
            try:
                # Large group message (group-encrypted)
                if raw.group_encrypted_content:
                    group_key = await self.get_group_key(raw.chat_id, raw.group_key_version)
                    plaintext = decrypt_with_group_key(
                        raw.group_encrypted_content, raw.group_nonce, group_key
                    )
                    messages.append(self._to_message(raw, plaintext))
                    continue

                # Version 2: Double Ratchet message
                if raw.ratchet_header and raw.version == 2:
                    plaintext = await self._decrypt_with_ratchet(raw)
                    messages.append(self._to_message(raw, plaintext))
                    continue

                # Version 1 (default): NaCl box
                sender = sender_cache.get(raw.from_agent)
                if not sender:
                    sender = await self._client.get_agent(raw.from_agent)
                    sender_cache[raw.from_agent] = sender
                plaintext = decrypt_from_sender(
                    raw.encrypted_content, raw.nonce,
                    sender.public_key, self._secret_key,
                )
                messages.append(self._to_message(raw, plaintext))
            except Exception as e:
                logger.warning(
                    "Failed to decrypt message %s from @%s: %s",
                    raw.id, raw.from_username, e,
                )
        return messages

    @staticmethod
    def _to_message(raw: RawIncomingMessage, plaintext: str) -> Message:
        return Message(
            id=raw.id,
            chat_id=raw.chat_id,
            from_agent=raw.from_agent,
            from_username=raw.from_username,
            from_name=raw.from_name,
            content=plaintext,
            content_type=raw.content_type,
            mentions=raw.mentions,
            reply_expected=bool(raw.reply_expected),
            reply_to=raw.reply_to,
            timestamp=raw.timestamp,
        )
