"""Crypto core for the X Chat platform adapter.

A thin, network-free wrapper around the ``chat_xdk`` binding. Everything
that touches the Chat XDK lives here so it can be unit-tested with a fake
``Chat`` object and so the adapter/API layers stay import-light. The SDK is
lazy-installed at first use via ``tools.lazy_deps`` (feature key
``platform.xchat``).

Responsibilities:

* key management        -> :meth:`XChatCrypto.load_keys` /
                           :meth:`XChatCrypto.generate_and_register_payload`
* session identity      -> :meth:`XChatCrypto.set_identity`
* signing-key roster    -> :meth:`XChatCrypto.set_signing_keys`
* message encryption    -> :meth:`XChatCrypto.encrypt_text` (with optional
                           attachments) and :meth:`XChatCrypto.encrypt_reply`
* event decryption      -> :meth:`XChatCrypto.decrypt_batch` (decrypt_events)
                           and :meth:`XChatCrypto.decrypt_one` (decrypt_event)
* media encryption      -> :meth:`XChatCrypto.encrypt_media` /
                           :meth:`XChatCrypto.decrypt_media` (stream cipher
                           under the conversation key)
* conversation keys     -> :meth:`XChatCrypto.prepare_conversation_key_change`
                           + :meth:`XChatCrypto.verify_key_binding` (initiate
                           brand-new conversations / rotate keys)

The decrypted-event dict shape follows the Chat XDK: ``{"type": "Message",
"id": ..., "sender_id": ..., "conversation_id": ..., "content": {"text":
...}}`` for messages, ``{"type": "KeyChange", ...}`` for conversation-key
rotations.
"""
from __future__ import annotations

import base64
import logging
from typing import Any, Optional

logger = logging.getLogger(__name__)


def _as_dict(obj: Any) -> dict[str, Any]:
    """Decrypted events come back as native objects; normalise to a dict."""
    if isinstance(obj, dict):
        return obj
    if hasattr(obj, "model_dump"):
        return obj.model_dump()
    try:
        return dict(obj)
    except Exception:
        return {}


def _load_chat_xdk():
    """Import (lazy-installing if needed) and return the ``chat_xdk`` module."""
    try:
        import chat_xdk  # type: ignore[import-not-found]
        return chat_xdk
    except ImportError:
        pass
    # Lazy-install path — same pattern as the telegram/matrix platform plugins.
    from tools.lazy_deps import ensure as _lazy_ensure

    _lazy_ensure("platform.xchat", prompt=False)
    import chat_xdk  # type: ignore[import-not-found]
    return chat_xdk


def _load_chat_class():
    """Import (lazy-installing if needed) and return ``chat_xdk.Chat``."""
    return _load_chat_xdk().Chat


def detect_mime_type(data: bytes) -> Optional[str]:
    """MIME sniff on PLAINTEXT bytes (Chat XDK helper)."""
    try:
        return _load_chat_xdk().detect_mime_type(bytes(data))
    except Exception:
        return None


def detect_image_dimensions(data: bytes) -> Optional[tuple[int, int]]:
    """(width, height) of PLAINTEXT image bytes, or None (Chat XDK helper)."""
    try:
        dims = _load_chat_xdk().detect_image_dimensions(bytes(data))
    except Exception:
        return None
    if not dims:
        return None
    try:
        # Bindings return either a (w, h) tuple or an object with attributes.
        if isinstance(dims, (tuple, list)):
            return int(dims[0]), int(dims[1])
        return int(dims.width), int(dims.height)
    except Exception:
        return None


class XChatCrypto:
    """Wraps a single unlocked ``chat_xdk.Chat`` instance for one bot identity."""

    def __init__(self, chat: Any = None) -> None:
        # ``chat`` injection keeps unit tests free of the native SDK.
        self.chat = chat if chat is not None else _load_chat_class()()
        self.signing_key_version: str = "1"
        self._identity_set = False

    # -- Key management -----------------------------------------------------

    def load_keys(self, private_keys_b64: str, signing_key_version: str = "1") -> None:
        """Import an existing private-key blob (from ``export_keys``) and adopt it.

        ``private_keys_b64`` is the base64 blob produced during registration
        (``hermes xchat setup``). Raises on a malformed blob.
        """
        blob = base64.b64decode(private_keys_b64.strip())
        self.chat.import_keys(blob, version=signing_key_version)
        self.signing_key_version = str(signing_key_version)

    def set_identity(self, user_id: str) -> None:
        """Set the session identity — every later encrypt call signs as this user."""
        self.chat.set_identity(str(user_id), self.signing_key_version)
        self._identity_set = True

    def set_cache_keys(self, enabled: bool = True) -> None:
        """Opt in to the SDK's verified conversation-key cache."""
        self.chat.set_cache_keys(enabled)

    def set_signing_keys(self, signing_keys: list[dict[str, str]]) -> None:
        """Replace the SDK's participant signing-key store (full roster each call)."""
        self.chat.set_signing_keys(signing_keys)

    def generate_and_register_payload(self) -> dict[str, Any]:
        """Generate fresh keypairs for a brand-new bot identity.

        Returns the registration body for ``POST /2/users/{id}/public_keys``
        plus the exported private-key blob (base64) to persist locally.
        Used by ``hermes xchat setup`` only — the adapter never generates keys.
        """
        reg = self.chat.generate_keypairs()
        version = str(reg.version) if getattr(reg, "version", None) is not None else "1"
        body = {
            "public_key": {
                "public_key": reg.public_key.public_key,
                "signing_public_key": reg.public_key.signing_public_key,
                "identity_public_key_signature": reg.public_key.identity_public_key_signature,
                "signing_public_key_signature": reg.public_key.signing_public_key_signature,
                "registration_method": reg.public_key.registration_method,
            },
            "version": version,
            "generate_version": bool(getattr(reg, "generate_version", False)),
        }
        exported = self.chat.export_keys()
        blob_b64 = base64.b64encode(bytes(exported)).decode("ascii") if exported else ""
        return {"registration": body, "version": version, "private_keys_b64": blob_b64}

    def verify_key_binding(
        self,
        identity_public_key_b64: str,
        signing_public_key_b64: str,
        identity_public_key_signature_b64: str,
    ) -> bool:
        """Verify a fetched public-key record's identity↔signing binding.

        MUST be called on every record before wrapping a conversation key to
        it (``prepare_conversation_key_change`` encrypts to whatever you pass
        — a substituted identity key would silently receive the key).
        """
        try:
            return bool(
                self.chat.verify_key_binding(
                    identity_public_key_b64,
                    signing_public_key_b64,
                    identity_public_key_signature_b64,
                )
            )
        except Exception:
            return False

    # -- Conversation-key setup (initiation / rotation) -----------------------

    def prepare_conversation_key_change(
        self,
        public_keys: list[dict[str, str]],
        *,
        conversation_id: Optional[str] = None,
    ) -> dict[str, Any]:
        """Generate + wrap a fresh conversation key for every participant.

        ``public_keys``: ``[{"user_id", "public_key", "key_version"}, ...]``
        (verified via :meth:`verify_key_binding` first). Returns the API body
        for ``POST /2/chat/conversations/{id}/keys`` plus the raw key:

        ``{"body": {...}, "conversation_key": bytes,
          "conversation_key_version": str}``
        """
        prepared = _as_dict(
            self.chat.prepare_conversation_key_change(
                public_keys, conversation_id=conversation_id
            )
        )
        body = {
            "conversation_key_version": prepared["conversation_key_version"],
            "conversation_participant_keys": [
                {
                    "user_id": pk["user_id"],
                    "encrypted_conversation_key": pk["encrypted_key"],
                    "public_key_version": pk["public_key_version"],
                }
                for pk in (prepared.get("participant_keys") or [])
            ],
            "action_signatures": [
                {
                    "message_id": sig["message_id"],
                    "encoded_message_event_detail": sig["encoded_message_event_detail"],
                    "message_event_signature": {
                        "signature": sig["signature"],
                        "public_key_version": sig["public_key_version"],
                        "signature_version": sig["signature_version"],
                    },
                }
                for sig in (prepared.get("action_signatures") or [])
            ],
        }
        return {
            "body": body,
            "conversation_key": prepared.get("conversation_key"),
            "conversation_key_version": str(prepared["conversation_key_version"]),
        }

    # -- Decryption ----------------------------------------------------------

    def decrypt_batch(self, events_b64: list[str]) -> dict[str, Any]:
        """Batch path — initial backlog load and KeyChange processing.

        ``decrypt_events`` extracts conversation keys from any KeyChange
        events in the batch (feeding the SDK's key cache when enabled), then
        decrypts every message. Signing keys come from the
        ``set_signing_keys`` store.

        NOTE: the events endpoint returns KeyChange events SEPARATELY in
        ``meta.conversation_key_events`` — callers must prepend those to the
        batch or no conversation key is ever extracted.
        """
        result = self.chat.decrypt_events(events_b64, None)
        messages = [
            {"event": _as_dict(m.get("event") if isinstance(m, dict) else m)}
            for m in (result.get("messages") or [])
        ]
        conv_keys = result.get("conversation_keys") or {}
        return {
            "messages": messages,
            "conversation_keys": conv_keys,
            "latest_key_version": (
                str(conv_keys.get("latest_version"))
                if conv_keys.get("latest_version") is not None
                else None
            ),
            "errors": result.get("errors") or {},
        }

    def decrypt_one(
        self, event_b64: str, conversation_keys: Optional[dict[str, bytes]] = None
    ) -> dict[str, Any]:
        """Single-event path — per-poll decryption with cached conversation keys."""
        return _as_dict(self.chat.decrypt_event(event_b64, conversation_keys, None))

    # -- Encryption ----------------------------------------------------------

    def encrypt_text(
        self,
        conversation_id: str,
        text: str,
        *,
        attachments: Optional[list[dict[str, Any]]] = None,
        conversation_key: Optional[bytes] = None,
        conversation_key_version: Optional[str] = None,
    ) -> dict[str, str]:
        """Encrypt + sign ``text`` (optionally with media attachments).

        Returns the X API send-message body. The conversation key is resolved
        from the SDK's verified-key cache (``set_cache_keys``) unless an
        explicit ``conversation_key`` + version pair is given (used right
        after key initiation, before any KeyChange event has been polled).
        Raises ``ValueError`` when no verified key is available.
        """
        kwargs: dict[str, Any] = {}
        if attachments:
            kwargs["attachments"] = attachments
        if conversation_key is not None:
            kwargs["conversation_key"] = conversation_key
            kwargs["conversation_key_version"] = conversation_key_version
        payload = self.chat.encrypt_message(str(conversation_id), text, **kwargs)
        return {
            "message_id": payload.message_id,
            "encoded_message_create_event": payload.encrypted_content,
            "encoded_message_event_signature": payload.encoded_event_signature,
        }

    def encrypt_reply(
        self,
        conversation_id: str,
        text: str,
        reply_to_event: dict[str, Any],
        *,
        attachments: Optional[list[dict[str, Any]]] = None,
    ) -> dict[str, str]:
        """Encrypt + sign ``text`` as a native threaded reply.

        ``reply_to_event`` is the DECRYPTED event dict of the message being
        replied to (from :meth:`decrypt_one` / :meth:`decrypt_batch`).
        Falls back to :meth:`encrypt_text` semantics on SDK versions without
        reply support.
        """
        kwargs: dict[str, Any] = {}
        if attachments:
            kwargs["attachments"] = attachments
        payload = self.chat.encrypt_reply(
            str(conversation_id), text, reply_to_event=reply_to_event, **kwargs
        )
        return {
            "message_id": payload.message_id,
            "encoded_message_create_event": payload.encrypted_content,
            "encoded_message_event_signature": payload.encoded_event_signature,
        }

    # -- Media (stream cipher under the conversation key) ---------------------

    def encrypt_media(self, plaintext: bytes, conversation_key: bytes) -> bytes:
        """Encrypt attachment bytes for upload (whole payload in memory)."""
        return bytes(self.chat.encrypt_stream(bytes(plaintext), conversation_key))

    def decrypt_media(self, ciphertext: bytes, conversation_key: bytes) -> bytes:
        """Decrypt a downloaded attachment blob.

        The key MUST be the conversation key for the *event's* key version —
        after a rotation, the latest key cannot decrypt older media.
        """
        return bytes(self.chat.decrypt_stream(bytes(ciphertext), conversation_key))


def message_text(event: dict[str, Any]) -> Optional[str]:
    """Pull the plain text out of a decrypted Message/MessageEdit event.

    Edits matter: the feed sometimes returns an edited message only as the
    edit event (the original is dropped), so skipping edits would make the
    message invisible to the bot forever.
    """
    if event.get("type") not in ("Message", "MessageEdit"):
        return None
    content = event.get("content") or {}
    if isinstance(content, dict):
        return content.get("text")
    return None


def message_attachments(event: dict[str, Any]) -> list[dict[str, Any]]:
    """Media attachment descriptors from a decrypted Message/MessageEdit event.

    Each entry carries ``media_hash_key`` (+ optional filename/width/height/
    filesize_bytes). Returns [] for non-message events or text-only messages.
    """
    if event.get("type") not in ("Message", "MessageEdit"):
        return []
    content = event.get("content") or {}
    if not isinstance(content, dict):
        return []
    raw = content.get("attachments") or event.get("attachments") or []
    out: list[dict[str, Any]] = []
    for att in raw if isinstance(raw, (list, tuple)) else []:
        att = _as_dict(att)
        if att.get("media_hash_key"):
            out.append(att)
    return out