"""High-level session classes wrapping crypto, auth, and the autogenerated API client."""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from typing import Any

import httpx

from arbi_client.api.user import logout as logout_api
from arbi_client.api.workspace import create_workspace as create_workspace_api
from arbi_client.client import AuthenticatedClient, Client
from arbi_client.models.task_update_message import TaskUpdateMessage
from arbi_client.models.workspace_create_request import WorkspaceCreateRequest
from arbi_client.sdk import auth as _auth
from arbi_client.sdk._types import QueryResult, SessionState, UserIdentity
from arbi_client.sdk.crypto import (
    decrypt_message,
    encrypt_message,
    encrypt_workspace_key_for_header,
    generate_workspace_key,
    unwrap_key_with_x25519,
)


def _unwrap_and_encrypt(wrapped_key_b64: str, encryption_private_key: bytes, session_public_key_b64: str) -> str:
    """Unwrap a workspace key (X25519 from DB) and SealedBox-encrypt with session public key.

    Called lazily — only when the encrypted header is actually needed (e.g. /open).
    """
    raw_key = unwrap_key_with_x25519(wrapped_key_b64, encryption_private_key)
    return encrypt_workspace_key_for_header(raw_key, session_public_key_b64)


# ---------------------------------------------------------------------------
# Shared helpers used by both ArbiSession and AsyncArbiSession
# ---------------------------------------------------------------------------


def _cache_workspace_header_impl(
    session: SessionState,
    workspace_ext_id: str,
    wrapped_key_b64: str | None,
    header: str | None,
) -> None:
    """Store a workspace key on the session state (no I/O)."""
    if header:
        session.workspace_headers[workspace_ext_id] = header
    elif wrapped_key_b64:
        session.wrapped_keys[workspace_ext_id] = wrapped_key_b64
    else:
        raise ValueError("Either wrapped_key_b64 or header must be provided")


def _workspace_key_header_impl(
    session: SessionState,
    identity: UserIdentity | None,
    active_workspace: str | None,
    workspace_ext_id: str | None,
) -> str:
    """Return a SealedBox-encrypted workspace key, lazily decrypting on first access (no I/O)."""
    ws_id = workspace_ext_id or active_workspace
    if ws_id is None:
        raise RuntimeError("No workspace selected — call select_workspace() first")

    header = session.workspace_headers.get(ws_id)
    if header is None:
        wrapped = session.wrapped_keys.get(ws_id)
        if wrapped is None:
            raise RuntimeError(f"No key for workspace {ws_id!r} — call cache_workspace_header() first")
        if identity is None:
            raise RuntimeError("Cannot unwrap workspace key without identity")
        header = _unwrap_and_encrypt(wrapped, identity.encryption_private_key, session.session_public_key_b64)
        session.workspace_headers[ws_id] = header

    return header


class ArbiSession:
    """Synchronous high-level ARBI session.

    Wraps key derivation, login, workspace key management, and provides
    an :class:`~arbi_client.client.AuthenticatedClient` for use with the
    autogenerated ``arbi_client.api.*`` functions.

    Usage::

        session = ArbiSession(base_url="https://app.arbi.city", deployment_domain="app.arbi.city")
        session.login(email="user@example.com", password="secret")  # pragma: allowlist secret
        ws = session.create_workspace("My Workspace")
        session.open_workspace(ws.external_id)  # gets workspace-scoped JWT

        # Use with autogenerated API functions
        from arbi_client.api.document import list_documents
        docs = list_documents.sync(client=session.authenticated_client)
    """

    def __init__(
        self,
        *,
        base_url: str,
        deployment_domain: str,
        timeout: float = 120.0,
        verify_ssl: bool = True,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._deployment_domain = deployment_domain
        self._timeout = timeout
        self._verify_ssl = verify_ssl

        self._base_client = Client(
            base_url=self._base_url,
            verify_ssl=self._verify_ssl,
            timeout=httpx.Timeout(self._timeout),
        )

        self._identity: UserIdentity | None = None
        self._session: SessionState | None = None
        self._auth_client: AuthenticatedClient | None = None
        self._active_workspace: str | None = None

    # ------------------------------------------------------------------
    # Auth
    # ------------------------------------------------------------------

    def register(
        self,
        *,
        email: str,
        password: str,
        given_name: str,
        family_name: str,
        verification_credential: str,
    ) -> UserIdentity:
        """Register a new user. Returns the :class:`UserIdentity` (not logged in yet)."""
        self._identity = _auth.register(
            self._base_client,
            email=email,
            password=password,
            given_name=given_name,
            family_name=family_name,
            deployment_domain=self._deployment_domain,
            verification_credential=verification_credential,
        )
        return self._identity

    def login(self, *, email: str, password: str) -> SessionState:
        """Derive keys from password, log in, and store session state.

        Automatically caches workspace headers from the login response
        so that :meth:`open_workspace` works without a separate
        :meth:`list_workspaces` call.
        """
        identity, session, workspaces = _auth.login(
            self._base_client,
            email=email,
            password=password,
            deployment_domain=self._deployment_domain,
        )
        self._identity = identity
        self._session = session
        self._auth_client = self._make_auth_client(session.access_token)
        self._cache_login_workspaces(workspaces)
        return session

    def login_with_identity(self, identity: UserIdentity) -> SessionState:
        """Log in with an existing :class:`UserIdentity` (skips Argon2 derivation)."""
        session, workspaces = _auth.login_with_identity(self._base_client, identity=identity)
        self._identity = identity
        self._session = session
        self._auth_client = self._make_auth_client(session.access_token)
        self._cache_login_workspaces(workspaces)
        return session

    def login_as_agent(self, *, agent_email: str, signing_key: bytes) -> SessionState:
        """Log in as a persistent agent using Ed25519 signing key.

        Works exactly like user login: derives X25519 encryption keys from the
        Ed25519 signing key, authenticates via signature, and caches workspace
        keys for subsequent ``open_workspace()`` calls.

        Args:
            agent_email: Agent's email address (e.g. ``agentname.usr-XXXXXXXX@deploymentdomain``).
            signing_key: Ed25519 private key bytes (64-byte key or 32-byte seed).
        """
        identity, session, workspaces = _auth.agent_login(
            self._base_client,
            agent_email=agent_email,
            signing_key=signing_key,
        )
        self._identity = identity
        self._session = session
        self._auth_client = self._make_auth_client(session.access_token)
        self._cache_login_workspaces(workspaces)
        return session

    def logout(self) -> None:
        """Log out and clear local state."""
        if self._auth_client is not None:
            logout_api.sync(client=self._auth_client)
        self._session = None
        self._auth_client = None
        self._active_workspace = None

    # ------------------------------------------------------------------
    # Workspace management
    # ------------------------------------------------------------------

    def select_workspace(self, workspace_ext_id: str) -> None:
        """Set the active workspace for subsequent operations."""
        self._active_workspace = workspace_ext_id

    def open_workspace(self, workspace_ext_id: str):
        """Open a workspace: store encrypted key in session, set active, return data.

        For user sessions: SealedBox-encrypts the workspace key with the session
        public key and sends it in the request body.

        SealedBox-encrypts the wrapped workspace key (from login response) with
        the session public key and sends it in the request body.

        Args:
            workspace_ext_id: Workspace external ID to open.

        Returns:
            :class:`WorkspaceOpenResponse` with ``workspace``,
            ``conversations``, ``documents``, and ``tags``.
        """
        from arbi_client.api.workspace import open_workspace as open_ws_api
        from arbi_client.models.workspace_open_request import WorkspaceOpenRequest

        self._require_session()

        # Agent sessions (no identity) have keys server-side; user sessions build the header
        ws_key_header = None
        if self._identity is not None:
            ws_key_header = self.workspace_key_header(workspace_ext_id)

        result = open_ws_api.sync(
            workspace_ext_id,
            client=self._auth_client,
            body=WorkspaceOpenRequest(workspace_key=ws_key_header),
        )
        if result is None or hasattr(result, "detail"):
            raise RuntimeError(f"Failed to open workspace: {result}")

        # Set active workspace (JWT stays the same — no new token)
        self._active_workspace = workspace_ext_id

        return result

    def cache_workspace_header(
        self,
        workspace_ext_id: str,
        wrapped_key_b64: str | None = None,
        header: str | None = None,
    ) -> None:
        """Store a workspace key for later use.

        Either stores a ``wrapped_key_b64`` (from DB, decrypted lazily on demand)
        or a pre-built ``header`` (already SealedBox-encrypted, e.g. from create_workspace).
        """
        self._require_session()
        _cache_workspace_header_impl(self._session, workspace_ext_id, wrapped_key_b64, header)

    def create_workspace(self, name: str, description: str = "") -> Any:
        """Create a workspace with a client-generated key.

        Generates a random 32-byte workspace key locally, encrypts it with the
        server's session public key, and sends it in the create request.

        Returns the workspace response object from the API.
        """
        self._require_session()

        # Client generates the workspace key
        workspace_key = generate_workspace_key()
        encrypted_key = encrypt_workspace_key_for_header(workspace_key, self._session.session_public_key_b64)

        workspace = create_workspace_api.sync(
            client=self._auth_client,
            body=WorkspaceCreateRequest(name=name, description=description, workspace_key=encrypted_key),
        )
        if workspace is None:
            raise RuntimeError("Workspace creation returned None")

        # Cache the header (for immediate use) and the DB-wrapped key (for re-wrapping to agents)
        header = encrypt_workspace_key_for_header(workspace_key, self._session.session_public_key_b64)
        self.cache_workspace_header(workspace.external_id, header=header)
        if workspace.wrapped_key:
            self._session.wrapped_keys[workspace.external_id] = workspace.wrapped_key
        self._active_workspace = workspace.external_id

        return workspace

    # ------------------------------------------------------------------
    # Header helpers
    # ------------------------------------------------------------------

    @property
    def access_token(self) -> str:
        """The current JWT access token."""
        self._require_session()
        return self._session.access_token

    @property
    def active_workspace(self) -> str | None:
        """The currently selected workspace external ID, or ``None``."""
        return self._active_workspace

    def workspace_key_header(self, workspace_ext_id: str | None = None) -> str:
        """Return a SealedBox-encrypted workspace key for /open requests.

        Lazily unwraps the X25519-wrapped key from DB and re-encrypts with
        the session public key on first access.

        Args:
            workspace_ext_id: Workspace to get the header for.
                Defaults to the active workspace.
        """
        self._require_session()
        return _workspace_key_header_impl(self._session, self._identity, self._active_workspace, workspace_ext_id)

    def headers(self) -> dict[str, str]:
        """Return a dict of common headers for raw API calls.

        After ``open_workspace()``, the Authorization header carries a
        workspace-scoped JWT — no separate ``workspace-key`` header needed.
        """
        self._require_session()
        return {
            "Authorization": f"Bearer {self._session.access_token}",
        }

    # ------------------------------------------------------------------
    # Authenticated client
    # ------------------------------------------------------------------

    @property
    def authenticated_client(self) -> AuthenticatedClient:
        """An :class:`AuthenticatedClient` with the Bearer token set."""
        self._require_session()
        return self._auth_client

    # ------------------------------------------------------------------
    # Document operations
    # ------------------------------------------------------------------

    def upload_documents(
        self,
        file_paths: list[str | Path],
        *,
        wait: bool = True,
        timeout: int = 300,
        shared: bool = True,
        config_ext_id: str | None = None,
        tag_ext_id: str | None = None,
        on_status: Callable[[TaskUpdateMessage], Any] | None = None,
    ) -> list[str]:
        """Upload files and optionally wait for processing via WebSocket.

        Args:
            file_paths: Local file paths to upload.
            wait: Block until processing completes (default ``True``).
            timeout: Max seconds to wait when *wait* is ``True``.
            shared: Share with workspace members.
            config_ext_id: Optional processing config.
            tag_ext_id: Optional tag to link documents to.
            on_status: Optional callback receiving a :class:`TaskUpdateMessage`
                for every status update via WebSocket while waiting.

        Returns:
            List of document external IDs in upload order.
        """
        from arbi_client.sdk.documents import upload_documents as _upload

        self._require_session()
        self._require_workspace()
        return _upload(
            client=self._auth_client,
            base_url=self._base_url,
            access_token=self._session.access_token,
            file_paths=file_paths,
            wait=wait,
            timeout=timeout,
            shared=shared,
            config_ext_id=config_ext_id,
            tag_ext_id=tag_ext_id,
            on_status=on_status,
        )

    def list_documents(self) -> list:
        """List all documents in the active workspace."""
        from arbi_client.api.document import list_documents as list_docs_api

        self._require_session()
        self._require_workspace()
        return list_docs_api.sync(client=self._auth_client) or []

    def delete_documents(self, doc_ext_ids: list[str]) -> Any:
        """Delete documents by their external IDs."""
        from arbi_client.sdk.documents import delete_documents as _delete

        self._require_session()
        return _delete(client=self._auth_client, doc_ext_ids=doc_ext_ids)

    # ------------------------------------------------------------------
    # Query / chat
    # ------------------------------------------------------------------

    def query(
        self,
        content: str,
        *,
        doc_ext_ids: list[str] | None = None,
        search_mode: str = "hybrid",
        tools: dict | None = None,
        config_ext_id: str | None = None,
        timeout: float = 600.0,
    ) -> str:
        """Run a query and return the answer string.

        If *tools* is not provided, a default ``retrieval_chunk`` tool is built
        from *doc_ext_ids* and *search_mode*.
        """
        from arbi_client.sdk.query import query as _query

        self._require_session()
        self._require_workspace()
        return _query(
            base_url=self._base_url,
            access_token=self._session.access_token,
            content=content,
            doc_ext_ids=doc_ext_ids,
            search_mode=search_mode,
            tools=tools,
            config_ext_id=config_ext_id,
            timeout=timeout,
            verify=self._verify_ssl,
        )

    def query_streaming(
        self,
        content: str,
        *,
        doc_ext_ids: list[str] | None = None,
        search_mode: str = "hybrid",
        tools: dict | None = None,
        config_ext_id: str | None = None,
        on_token: Callable[[str], Any] | None = None,
        on_agent_step: Callable[[dict], Any] | None = None,
        timeout: float = 600.0,
    ) -> QueryResult:
        """Run a streaming query with optional per-event callbacks.

        Returns a :class:`QueryResult` with ``content``, ``events``,
        ``metadata``, and ``assistant_message_ext_id``.
        """
        from arbi_client.sdk.query import query_streaming as _query_streaming

        self._require_session()
        self._require_workspace()
        return _query_streaming(
            base_url=self._base_url,
            access_token=self._session.access_token,
            content=content,
            doc_ext_ids=doc_ext_ids,
            search_mode=search_mode,
            tools=tools,
            config_ext_id=config_ext_id,
            on_token=on_token,
            on_agent_step=on_agent_step,
            timeout=timeout,
            verify=self._verify_ssl,
        )

    # ------------------------------------------------------------------
    # Workspace listing & sharing
    # ------------------------------------------------------------------

    def list_workspaces(self) -> list:
        """List all workspaces the user has access to.

        Stores wrapped keys lazily — decryption only happens when a workspace is opened.
        """
        from arbi_client.api.user import get_user_workspaces as get_ws_api

        self._require_session()
        workspaces = get_ws_api.sync(client=self._auth_client)
        if workspaces is None:
            return []
        for ws in workspaces:
            wrapped = getattr(ws, "wrapped_key", None)
            if wrapped and ws.external_id not in self._session.wrapped_keys:
                self._session.wrapped_keys[ws.external_id] = wrapped
        return workspaces

    def add_workspace_users(
        self,
        emails: list[str],
        role: str = "collaborator",
    ) -> list:
        """Add users to the active workspace.

        Args:
            emails: Email addresses of users to invite.
            role: Role to assign (``"owner"``, ``"collaborator"``, ``"guest"``).

        Returns:
            List of ``WorkspaceUserResponse`` objects.
        """
        from arbi_client.api.workspace import add_workspace_users as add_users_api
        from arbi_client.models.add_workspace_users_request import AddWorkspaceUsersRequest
        from arbi_client.models.workspace_role import WorkspaceRole

        self._require_session()
        self._require_workspace()
        return (
            add_users_api.sync(
                client=self._auth_client,
                body=AddWorkspaceUsersRequest(emails=emails, role=WorkspaceRole(role)),
            )
            or []
        )

    def accept_workspace_invite(
        self,
        workspace_ext_id: str,
        wrapped_key_b64: str,
    ) -> None:
        """Accept a workspace invitation by unwrapping the shared key.

        After calling this, the workspace is selected as the active workspace.
        """
        self.cache_workspace_header(workspace_ext_id, wrapped_key_b64=wrapped_key_b64)
        self.select_workspace(workspace_ext_id)

    # ------------------------------------------------------------------
    # Tag operations
    # ------------------------------------------------------------------

    def create_tag(self, name: str) -> Any:
        """Create a tag in the active workspace."""
        from arbi_client.api.tag import create_tag as create_tag_api
        from arbi_client.models.create_tag_request import CreateTagRequest

        self._require_session()
        self._require_workspace()
        return create_tag_api.sync(client=self._auth_client, body=CreateTagRequest(name=name))

    def list_tags(self) -> list:
        """List all tags in the active workspace."""
        from arbi_client.api.tag import list_tags as list_tags_api

        self._require_session()
        self._require_workspace()
        return list_tags_api.sync(client=self._auth_client) or []

    def tag_documents(self, tag_ext_id: str, doc_ext_ids: list[str]) -> list:
        """Link documents to a tag by creating doc-tag associations.

        Args:
            tag_ext_id: The tag to associate with.
            doc_ext_ids: Document external IDs to tag.

        Returns:
            List of :class:`DocTagResponse` objects.
        """
        from arbi_client.api.document import create_doctags as create_doctags_api
        from arbi_client.models.create_doc_tag_request import CreateDocTagRequest

        self._require_session()
        self._require_workspace()
        return (
            create_doctags_api.sync(
                client=self._auth_client,
                body=CreateDocTagRequest(tag_ext_id=tag_ext_id, doc_ext_ids=doc_ext_ids),
            )
            or []
        )

    # ------------------------------------------------------------------
    # Conversation operations
    # ------------------------------------------------------------------

    def list_conversations(self) -> list:
        """List all conversations in the active workspace."""
        import types

        from arbi_client.api.conversation import list_conversations as list_convos_api

        self._require_session()
        self._require_workspace()
        result = list_convos_api.sync(client=self._auth_client) or []
        # Endpoint returns list[Any] (untyped) — convert dicts to SimpleNamespace for attribute access
        return [types.SimpleNamespace(**item) if isinstance(item, dict) else item for item in result]

    def delete_conversation(self, convo_ext_id: str) -> Any:
        """Delete a conversation and all its messages."""
        from arbi_client.api.conversation import (
            delete_conversation as delete_convo_api,
        )

        self._require_session()
        return delete_convo_api.sync(convo_ext_id, client=self._auth_client)

    # ------------------------------------------------------------------
    # Config operations
    # ------------------------------------------------------------------

    def save_config(self, **kwargs: Any) -> Any:
        """Save a configuration.

        Pass any fields accepted by ``ConfigUpdateData`` as keyword arguments,
        e.g. ``session.save_config(title="My Config")``.
        """
        from arbi_client.api.configs import save_config as save_config_api
        from arbi_client.models.config_update_data import ConfigUpdateData

        self._require_session()
        result = save_config_api.sync(
            client=self._auth_client,
            body=ConfigUpdateData.from_dict(kwargs),
        )
        if result is None or hasattr(result, "detail"):
            raise RuntimeError(f"Save config failed: {result}")
        return result

    def delete_config(self, config_ext_id: str) -> Any:
        """Delete a configuration by its external ID."""
        from arbi_client.api.configs import delete_config as delete_config_api

        self._require_session()
        return delete_config_api.sync(config_ext_id, client=self._auth_client)

    # ------------------------------------------------------------------
    # E2E messaging
    # ------------------------------------------------------------------

    def encrypt_message(self, plaintext: str, recipient_pub_key_b64: str) -> str:
        """Encrypt a message for a recipient using X25519 DH."""
        self._require_identity()
        return encrypt_message(plaintext, recipient_pub_key_b64, self._identity.encryption_private_key)

    def decrypt_message(self, encrypted_b64: str, sender_pub_key_b64: str) -> str:
        """Decrypt a message from a sender using X25519 DH."""
        self._require_identity()
        return decrypt_message(encrypted_b64, sender_pub_key_b64, self._identity.encryption_private_key)

    # ------------------------------------------------------------------
    # Identity / session access
    # ------------------------------------------------------------------

    @property
    def identity(self) -> UserIdentity:
        """The current user's identity."""
        self._require_identity()
        return self._identity

    @property
    def session_state(self) -> SessionState:
        """The current session state."""
        self._require_session()
        return self._session

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    def close(self) -> None:
        """Release resources (does **not** call logout)."""
        pass

    def __enter__(self) -> ArbiSession:
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _make_auth_client(self, token: str) -> AuthenticatedClient:
        return AuthenticatedClient(
            base_url=self._base_url,
            token=token,
            verify_ssl=self._verify_ssl,
            timeout=httpx.Timeout(self._timeout),
        )

    def _cache_login_workspaces(self, workspaces) -> None:
        """Store wrapped keys from login response (decrypted lazily on demand)."""
        for ws in workspaces:
            if ws.wrapped_key and ws.external_id not in self._session.wrapped_keys:
                self._session.wrapped_keys[ws.external_id] = ws.wrapped_key

    def _require_identity(self) -> None:
        if self._identity is None:
            raise RuntimeError("Not logged in — call login() or register() first")

    def _require_session(self) -> None:
        if self._session is None:
            raise RuntimeError("Not logged in — call login() first")

    def _require_workspace(self) -> None:
        if self._active_workspace is None:
            raise RuntimeError("No workspace selected — call select_workspace() or create_workspace() first")


class AsyncArbiSession:
    """Async variant of :class:`ArbiSession`.

    All mutating methods are ``async``. Properties remain synchronous.
    """

    def __init__(
        self,
        *,
        base_url: str,
        deployment_domain: str,
        timeout: float = 120.0,
        verify_ssl: bool = True,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._deployment_domain = deployment_domain
        self._timeout = timeout
        self._verify_ssl = verify_ssl

        self._base_client = Client(
            base_url=self._base_url,
            verify_ssl=self._verify_ssl,
            timeout=httpx.Timeout(self._timeout),
        )

        self._identity: UserIdentity | None = None
        self._session: SessionState | None = None
        self._auth_client: AuthenticatedClient | None = None
        self._active_workspace: str | None = None

    # ------------------------------------------------------------------
    # Auth
    # ------------------------------------------------------------------

    async def register(
        self,
        *,
        email: str,
        password: str,
        given_name: str,
        family_name: str,
        verification_credential: str,
    ) -> UserIdentity:
        self._identity = await _auth.aregister(
            self._base_client,
            email=email,
            password=password,
            given_name=given_name,
            family_name=family_name,
            deployment_domain=self._deployment_domain,
            verification_credential=verification_credential,
        )
        return self._identity

    async def login(self, *, email: str, password: str) -> SessionState:
        identity, session, workspaces = await _auth.alogin(
            self._base_client,
            email=email,
            password=password,
            deployment_domain=self._deployment_domain,
        )
        self._identity = identity
        self._session = session
        self._auth_client = self._make_auth_client(session.access_token)
        self._cache_login_workspaces(workspaces)
        return session

    async def login_with_identity(self, identity: UserIdentity) -> SessionState:
        session, workspaces = await _auth.alogin_with_identity(self._base_client, identity=identity)
        self._identity = identity
        self._session = session
        self._auth_client = self._make_auth_client(session.access_token)
        self._cache_login_workspaces(workspaces)
        return session

    async def login_as_agent(self, *, agent_email: str, signing_key: bytes) -> SessionState:
        """Async variant of :meth:`ArbiSession.login_as_agent`."""
        identity, session, workspaces = await _auth.aagent_login(
            self._base_client,
            agent_email=agent_email,
            signing_key=signing_key,
        )
        self._identity = identity
        self._session = session
        self._auth_client = self._make_auth_client(session.access_token)
        self._cache_login_workspaces(workspaces)
        return session

    async def logout(self) -> None:
        if self._auth_client is not None:
            await logout_api.asyncio(client=self._auth_client)
        self._session = None
        self._auth_client = None
        self._active_workspace = None

    # ------------------------------------------------------------------
    # Workspace management
    # ------------------------------------------------------------------

    def select_workspace(self, workspace_ext_id: str) -> None:
        self._active_workspace = workspace_ext_id

    async def open_workspace(self, workspace_ext_id: str):
        """Open a workspace (async variant of :meth:`ArbiSession.open_workspace`)."""
        from arbi_client.api.workspace import open_workspace as open_ws_api
        from arbi_client.models.workspace_open_request import WorkspaceOpenRequest

        self._require_session()

        ws_key_header = None
        if self._identity is not None:
            ws_key_header = self.workspace_key_header(workspace_ext_id)

        result = await open_ws_api.asyncio(
            workspace_ext_id,
            client=self._auth_client,
            body=WorkspaceOpenRequest(workspace_key=ws_key_header),
        )
        if result is None or hasattr(result, "detail"):
            raise RuntimeError(f"Failed to open workspace: {result}")

        # Set active workspace (JWT stays the same — no new token)
        self._active_workspace = workspace_ext_id

        return result

    def cache_workspace_header(
        self,
        workspace_ext_id: str,
        wrapped_key_b64: str | None = None,
        header: str | None = None,
    ) -> None:
        """Store a workspace key for later use (async variant).

        Either stores a ``wrapped_key_b64`` (from DB, decrypted lazily on demand)
        or a pre-built ``header`` (already SealedBox-encrypted).
        """
        self._require_session()
        _cache_workspace_header_impl(self._session, workspace_ext_id, wrapped_key_b64, header)

    async def create_workspace(self, name: str, description: str = "") -> Any:
        """Create a workspace with a client-generated key (async variant)."""
        from arbi_client.api.workspace import create_workspace as cw_api

        self._require_session()

        # Client generates the workspace key
        workspace_key = generate_workspace_key()
        encrypted_key = encrypt_workspace_key_for_header(workspace_key, self._session.session_public_key_b64)

        workspace = await cw_api.asyncio(
            client=self._auth_client,
            body=WorkspaceCreateRequest(name=name, description=description, workspace_key=encrypted_key),
        )
        if workspace is None:
            raise RuntimeError("Workspace creation returned None")

        # Cache the header (for immediate use) and the DB-wrapped key (for re-wrapping to agents)
        header = encrypt_workspace_key_for_header(workspace_key, self._session.session_public_key_b64)
        self.cache_workspace_header(workspace.external_id, header=header)
        if workspace.wrapped_key:
            self._session.wrapped_keys[workspace.external_id] = workspace.wrapped_key
        self._active_workspace = workspace.external_id

        return workspace

    # ------------------------------------------------------------------
    # Header helpers (same as sync — no I/O)
    # ------------------------------------------------------------------

    @property
    def access_token(self) -> str:
        self._require_session()
        return self._session.access_token

    @property
    def active_workspace(self) -> str | None:
        return self._active_workspace

    def workspace_key_header(self, workspace_ext_id: str | None = None) -> str:
        """Return a SealedBox-encrypted workspace key, lazily decrypting on first access."""
        self._require_session()
        return _workspace_key_header_impl(self._session, self._identity, self._active_workspace, workspace_ext_id)

    def headers(self) -> dict[str, str]:
        """Return a dict of common headers for raw API calls.

        After ``open_workspace()``, the Authorization header carries a
        workspace-scoped JWT — no separate ``workspace-key`` header needed.
        """
        self._require_session()
        return {
            "Authorization": f"Bearer {self._session.access_token}",
        }

    @property
    def authenticated_client(self) -> AuthenticatedClient:
        self._require_session()
        return self._auth_client

    # ------------------------------------------------------------------
    # E2E messaging
    # ------------------------------------------------------------------

    def encrypt_message(self, plaintext: str, recipient_pub_key_b64: str) -> str:
        self._require_identity()
        return encrypt_message(plaintext, recipient_pub_key_b64, self._identity.encryption_private_key)

    def decrypt_message(self, encrypted_b64: str, sender_pub_key_b64: str) -> str:
        self._require_identity()
        return decrypt_message(encrypted_b64, sender_pub_key_b64, self._identity.encryption_private_key)

    # ------------------------------------------------------------------
    # Identity / session access
    # ------------------------------------------------------------------

    @property
    def identity(self) -> UserIdentity:
        self._require_identity()
        return self._identity

    @property
    def session_state(self) -> SessionState:
        self._require_session()
        return self._session

    # ------------------------------------------------------------------
    # Context manager
    # ------------------------------------------------------------------

    async def close(self) -> None:
        pass

    async def __aenter__(self) -> AsyncArbiSession:
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.close()

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _make_auth_client(self, token: str) -> AuthenticatedClient:
        return AuthenticatedClient(
            base_url=self._base_url,
            token=token,
            verify_ssl=self._verify_ssl,
            timeout=httpx.Timeout(self._timeout),
        )

    def _cache_login_workspaces(self, workspaces) -> None:
        """Store wrapped keys from login response (decrypted lazily on demand)."""
        for ws in workspaces:
            if ws.wrapped_key and ws.external_id not in self._session.wrapped_keys:
                self._session.wrapped_keys[ws.external_id] = ws.wrapped_key

    def _require_identity(self) -> None:
        if self._identity is None:
            raise RuntimeError("Not logged in — call login() or register() first")

    def _require_session(self) -> None:
        if self._session is None:
            raise RuntimeError("Not logged in — call login() first")
