simplex-chat/packages/simplex-chat-python/src/simplex_chat/core.py

"""Typed async wrapper around libsimplex's C ABI functions.

Users interact with `Bot` / `ChatApi`. This module is exposed as
`simplex_chat.core` for tests, the api.ChatApi class and image resizing
(`resize_image`, `chat_resize_image`, `chat_resize_image_uri`).
"""

from __future__ import annotations

import asyncio
import base64
import ctypes
import json
from concurrent.futures import Executor
from enum import StrEnum
from typing import Any, Literal, TypedDict, overload

from . import _native
from .types import CR, CEvt, T


class ChatError(Exception):
    """Base class for every failure of a chat command.

    Catch this for both `ChatAPIError` and `api.ChatCommandError`.
    """


class ChatAPIError(ChatError):
    """Raised when chat_send_cmd / chat_recv_msg_wait returns a chat error."""

    def __init__(self, message: str, chat_error: T.ChatError | None = None):
        super().__init__(message)
        self.chat_error = chat_error

    @property
    def error_type(self) -> str | None:
        """Tag of the nested `errorType`, e.g. `noActiveUser`, or None."""
        return self._nested("errorType").get("type")

    @property
    def store_error_type(self) -> str | None:
        """Tag of the nested `storeError`, e.g. `duplicateName`, or None."""
        return self._nested("storeError").get("type")

    @property
    def command_error(self) -> str | None:
        """What the core says the caller did wrong, or None.

        The only part of a `commandError` worth reading: the tag says nothing.
        """
        error = self._nested("errorType")
        return error.get("message") if error.get("type") == "commandError" else None

    def _nested(self, key: str) -> dict[str, Any]:
        nested = (self.chat_error or {}).get(key)  # type: ignore[attr-defined]
        return nested if isinstance(nested, dict) else {}


class ChatInitError(Exception):
    """Raised when chat_migrate_init returns a DBMigrationResult error."""

    def __init__(self, message: str, db_migration_error: dict[str, Any]):
        super().__init__(message)
        self.db_migration_error = db_migration_error


class MigrationConfirmation(StrEnum):
    YES_UP = "yesUp"
    YES_UP_DOWN = "yesUpDown"
    CONSOLE = "console"
    ERROR = "error"


class ImageSize(StrEnum):
    IMAGE = "image"
    PROFILE_PREVIEW = "profilePreview"
    MESSAGE_PREVIEW = "messagePreview"


class CryptoArgs(TypedDict):  # wire-format JSON; camelCase fields
    fileKey: str
    fileNonce: str


def _read_and_free(ptr: int | None) -> str:
    """Copy a Haskell-allocated null-terminated UTF-8 string and free its buffer.

    Mirrors HandleCResult in packages/simplex-chat-nodejs/cpp/simplex.cc:157-165.
    """
    if not ptr:
        raise RuntimeError("null pointer returned from libsimplex")
    try:
        return ctypes.string_at(ptr).decode("utf-8")
    finally:
        _native.libc().free(ctypes.c_void_p(ptr))


def _check_c_int(name: str, value: int) -> None:
    # ctypes silently wraps ints that do not fit C int.
    if ctypes.c_int(value).value != value:
        raise ValueError(f"{name} {value} does not fit C int")


async def chat_send_cmd(ctrl: int, cmd: str) -> CR.ChatResponse:
    def _call() -> str:
        ptr = _native.lib().chat_send_cmd(ctrl, cmd.encode("utf-8"))
        return _read_and_free(ptr)

    raw = await asyncio.to_thread(_call)
    parsed = json.loads(raw)
    if "result" in parsed and isinstance(parsed["result"], dict):
        return parsed["result"]  # type: ignore[return-value]
    err = parsed.get("error")
    if isinstance(err, dict):
        raise ChatAPIError(f"chat command error: {err.get('type')}", err)  # type: ignore[arg-type]
    raise ChatAPIError(f"invalid chat command result: {raw[:200]}")


async def chat_recv_msg_wait(
    ctrl: int, wait_us: int = 500_000, executor: Executor | None = None
) -> CEvt.ChatEvent | None:
    def _call() -> str:
        # On timeout, the C side returns a non-NULL pointer to a single NUL byte
        # (see Mobile.hs `fromMaybe ""`), so `_read_and_free` returns "" — no
        # NULL-pointer guard is needed here.
        ptr = _native.lib().chat_recv_msg_wait(ctrl, wait_us)
        return _read_and_free(ptr)

    if executor is None:
        raw = await asyncio.to_thread(_call)
    else:
        raw = await asyncio.get_running_loop().run_in_executor(executor, _call)
    if not raw:
        return None
    parsed = json.loads(raw)
    if "result" in parsed and isinstance(parsed["result"], dict):
        return parsed["result"]  # type: ignore[return-value]
    err = parsed.get("error")
    if isinstance(err, dict):
        raise ChatAPIError(f"chat event error: {err.get('type')}", err)  # type: ignore[arg-type]
    raise ChatAPIError(f"invalid chat event: {raw[:200]}")


async def chat_migrate_init(
    db_path: str,
    db_key: str,
    confirm: MigrationConfirmation,
    queue_size: int | None = None,
) -> int:
    """Initialize chat controller. Returns opaque ctrl pointer as Python int.

    `queue_size` is the size of internal queues; the core default is used when None.
    """
    if queue_size is not None:
        _check_c_int("queue_size", queue_size)

    init_queue = _native.migrate_init_queue() if queue_size is not None else None

    def _call() -> tuple[int, str]:
        ctrl = ctypes.c_void_p()
        args = (db_path.encode("utf-8"), db_key.encode("utf-8"), confirm.encode("utf-8"))
        if init_queue is None:
            ptr = _native.lib().chat_migrate_init(*args, ctypes.byref(ctrl))
        else:
            ptr = init_queue(*args, queue_size, ctypes.byref(ctrl))
        return (ctrl.value or 0, _read_and_free(ptr))

    ctrl_val, raw = await asyncio.to_thread(_call)
    parsed = json.loads(raw)
    if parsed.get("type") == "ok":
        if not ctrl_val:
            # ABI invariant: type=="ok" → out-param written. Defensive guard so a
            # broken libsimplex doesn't hand us a NULL controller that would only
            # crash on first use much later.
            raise RuntimeError("chat_migrate_init returned ok but did not set ctrl pointer")
        return ctrl_val
    raise ChatInitError(
        "Database or migration error (see db_migration_error)",
        parsed,
    )


async def chat_close_store(ctrl: int) -> None:
    def _call() -> str:
        ptr = _native.lib().chat_close_store(ctrl)
        return _read_and_free(ptr)

    res = await asyncio.to_thread(_call)
    if res:
        raise RuntimeError(res)


async def chat_write_file(ctrl: int, path: str, data: bytes) -> CryptoArgs:
    def _call() -> str:
        ptr = _native.lib().chat_write_file(ctrl, path.encode("utf-8"), data, len(data))
        return _read_and_free(ptr)

    raw = await asyncio.to_thread(_call)
    return _crypto_args_result(raw)


async def chat_read_file(path: str, args: CryptoArgs) -> bytes:
    def _call() -> bytes:
        ptr = _native.lib().chat_read_file(
            path.encode("utf-8"),
            args["fileKey"].encode("utf-8"),
            args["fileNonce"].encode("utf-8"),
        )
        if not ptr:
            raise RuntimeError("chat_read_file returned null")
        addr = ctypes.cast(ptr, ctypes.c_void_p).value
        assert addr is not None  # `if not ptr` above already filtered NULL
        try:
            status = ctypes.cast(addr, ctypes.POINTER(ctypes.c_uint8))[0]
            if status == 1:
                msg = ctypes.string_at(addr + 1).decode("utf-8")
                raise RuntimeError(msg)
            if status != 0:
                raise RuntimeError(f"unexpected status {status} from chat_read_file")
            # `addr + 1` is unaligned for a uint32 read. On the supported platforms
            # (linux-x86_64, linux-aarch64, macos-aarch64, windows-x86_64) this is
            # silently handled; matches the Node.js binding (cpp/simplex.cc:344).
            length = ctypes.cast(addr + 1, ctypes.POINTER(ctypes.c_uint32))[0]
            return ctypes.string_at(addr + 5, length)
        finally:
            _native.libc().free(ctypes.c_void_p(addr))

    return await asyncio.to_thread(_call)


async def chat_encrypt_file(ctrl: int, src: str, dst: str) -> CryptoArgs:
    def _call() -> str:
        ptr = _native.lib().chat_encrypt_file(ctrl, src.encode("utf-8"), dst.encode("utf-8"))
        return _read_and_free(ptr)

    return _crypto_args_result(await asyncio.to_thread(_call))


async def chat_decrypt_file(src: str, args: CryptoArgs, dst: str) -> None:
    def _call() -> str:
        ptr = _native.lib().chat_decrypt_file(
            src.encode("utf-8"),
            args["fileKey"].encode("utf-8"),
            args["fileNonce"].encode("utf-8"),
            dst.encode("utf-8"),
        )
        return _read_and_free(ptr)

    res = await asyncio.to_thread(_call)
    if res:
        raise RuntimeError(res)


def _crypto_args_result(raw: str) -> CryptoArgs:
    parsed = json.loads(raw)
    if parsed.get("type") == "result":
        return parsed["cryptoArgs"]
    if parsed.get("type") == "error":
        raise RuntimeError(parsed.get("writeError", "unknown write error"))
    raise RuntimeError(f"unexpected result: {raw[:200]}")


async def chat_resize_image_uri(data: bytes, max_size: int) -> str:
    _check_c_int("len(data)", len(data))
    _check_c_int("max_size", max_size)
    resize = _native.resize_image_uri()

    def _call() -> str:
        return _read_and_free(resize(data, len(data), max_size))

    return _resized_image(await asyncio.to_thread(_call))


async def chat_resize_image(
    data: bytes, max_size: int, image_format: Literal["jpeg", "png"]
) -> bytes:
    return await _resize_image_data(data, max_size, 1 if image_format == "png" else 0)


@overload
async def resize_image(data: bytes, size: Literal[ImageSize.IMAGE]) -> bytes: ...


@overload
async def resize_image(
    data: bytes, size: Literal[ImageSize.PROFILE_PREVIEW, ImageSize.MESSAGE_PREVIEW]
) -> str: ...


@overload
async def resize_image(data: bytes, size: ImageSize) -> bytes | str: ...


async def resize_image(data: bytes, size: ImageSize) -> bytes | str:
    """Resize image to the app size limit: JPEG or PNG for an image, a data URI for a preview."""
    match size:
        case ImageSize.IMAGE:
            return await _resize_image_data(data, 261120, -1)
        case ImageSize.PROFILE_PREVIEW:
            return await chat_resize_image_uri(data, 12500)
        case ImageSize.MESSAGE_PREVIEW:
            return await chat_resize_image_uri(data, 14000)


async def _resize_image_data(data: bytes, max_size: int, png: int) -> bytes:
    _check_c_int("len(data)", len(data))
    _check_c_int("max_size", max_size)
    resize = _native.resize_image()

    def _call() -> str:
        return _read_and_free(resize(data, len(data), max_size, png))

    return base64.b64decode(_resized_image(await asyncio.to_thread(_call)))


def _resized_image(res: str) -> str:
    if not res:
        raise RuntimeError("image cannot be decoded or does not fit max_size")
    return res