# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import re
import time
from collections import defaultdict
from collections.abc import Mapping
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Any,
    Literal,
    TypeVar,
    Union,
    cast,
)
from uuid import uuid4

import msgspec

from marimo import _loggers
from marimo._ast.app_config import _AppConfig
from marimo._config.config import MarimoConfig
from marimo._data.models import DataTableSource
from marimo._messaging.notebook.document import NotebookCell
from marimo._messaging.notebook.outputs import CellOutputs
from marimo._types.encodable import Encodable
from marimo._types.ids import CellId_t, RequestId, UIElementId, WidgetModelId

LOGGER = _loggers.marimo_logger()

if TYPE_CHECKING:
    from collections.abc import Iterator

    from starlette.datastructures import URL
    from starlette.requests import HTTPConnection


def kebab_case(name: str) -> str:
    """Convert PascalCase to kebab-case.

    Removes 'Command' suffix and converts to kebab-case for discriminated union tags.
    Handles acronyms by keeping consecutive uppercase letters together.
    """
    name = name.removesuffix("Command")  # Remove 'Command' (7 characters)
    if not name:
        return name
    # Insert hyphens before uppercase letters that follow lowercase letters
    # or before an uppercase letter followed by a lowercase letter (handling acronyms)
    # This groups consecutive uppercase letters together (e.g., "SQL" -> "sql")
    s1 = re.sub("(.)([A-Z][a-z]+)", r"\1-\2", name)
    s2 = re.sub("([a-z0-9])([A-Z])", r"\1-\2", s1)
    return s2.lower()


class Command(
    msgspec.Struct, rename="camel", tag_field="type", tag=kebab_case
):
    """Base class for runtime commands.

    Discriminated union base using msgspec for serialization. Command subclasses
    are automatically tagged with their kebab-case name and serialized with
    camelCase field names. The "type" tag discriminates between commands during
    deserialization for type-safe routing.
    """


T = TypeVar("T")
# Union required: mypy rejects `T | list[T]` as a generic type alias target.
ListOrValue = Union[T, list[T]]
SerializedQueryParams = dict[str, ListOrValue[str]]
Primitive = str | bool | int | float
SerializedCLIArgs = dict[str, ListOrValue[Primitive]]


def _user_to_dict(user: Any) -> dict[str, Encodable]:
    """Normalize an auth user into a serializable dict.

    Starlette's authentication middleware sets `request.scope["user"]` to a
    `BaseUser` instance (commonly `SimpleUser`). Storing the raw object on
    `HTTPRequest.user` breaks msgspec serialization on the IPC path. Convert
    to a dict matching the documented contract: "info from authentication
    middleware (e.g., is_authenticated, username)".
    """
    if user is None:
        return {}
    if isinstance(user, dict):
        return user
    return {
        "username": getattr(user, "username", ""),
        "is_authenticated": getattr(user, "is_authenticated", False),
        "display_name": getattr(user, "display_name", ""),
    }


def _meta_to_dict(meta: Any) -> dict[str, Encodable]:
    """Normalize user-defined `request.scope["meta"]` at the boundary.

    Trusted shape per the public contract (`mo.app_meta().request`); we
    just enforce dict-ness at the seam. Non-`Encodable` values inside will
    surface at IPC encode time rather than being silently coerced.
    """
    if meta is None or not isinstance(meta, dict):
        return {}
    return cast("dict[str, Encodable]", meta)


@dataclass
class HTTPRequest(Mapping[str, Any]):
    """Serializable HTTP request representation.

    Mimics Starlette/FastAPI Request but is pickle-able and contains only a safe
    subset of data. Excludes session and auth to prevent exposing sensitive data.

    Attributes:
        url: Serialized URL with path, port, scheme, netloc, query, hostname.
        base_url: Serialized base URL.
        headers: Request headers (marimo-specific headers excluded).
        query_params: Query parameters mapped to lists of values.
        path_params: Path parameters from the URL route.
        cookies: Request cookies.
        meta: User-defined storage for custom data.
        user: User info from authentication middleware (e.g., is_authenticated, username).
    """

    url: dict[str, Encodable]  # Serialized URL
    base_url: dict[str, Encodable]  # Serialized URL
    headers: dict[str, str]  # Raw headers
    query_params: dict[str, list[str]]  # Raw query params
    path_params: dict[str, Encodable]
    cookies: dict[str, str]
    meta: dict[str, Encodable]  # User-defined storage
    user: Encodable

    # We don't include session or auth because they may contain
    # information that the app author does not want to expose.

    # session: dict[str, Any]
    # auth: Any

    def __getitem__(self, key: str) -> Any:
        return self.__dict__[key]

    def __iter__(self) -> Iterator[str]:
        return iter(self.__dict__)

    def __len__(self) -> int:
        return len(self.__dict__)

    def _display_(self) -> Any:
        try:
            import dataclasses

            return dataclasses.asdict(self)
        except TypeError:
            return self.__dict__

    def __repr__(self) -> str:
        return f"HTTPRequest(path={self.url['path']!r}, params={len(self.query_params)})"

    @staticmethod
    def from_request(request: HTTPConnection) -> HTTPRequest:
        def _url_to_dict(url: URL) -> dict[str, Encodable]:
            return {
                "path": url.path,
                "port": url.port,
                "scheme": url.scheme,
                "netloc": url.netloc,
                "query": url.query,
                "hostname": url.hostname,
            }

        # Convert URL to dict
        url_dict = _url_to_dict(request.url)

        # Convert base_url to dict
        base_url_dict = _url_to_dict(request.base_url)

        # Convert query params to dict[str, list[str]]
        query_params: dict[str, list[str]] = defaultdict(list)
        for k, v in request.query_params.multi_items():
            query_params[k].append(str(v))

        # Convert headers to dict, remove all marimo-specific headers
        headers: dict[str, str] = {}
        for k, v in request.headers.items():
            if not k.startswith(("marimo", "x-marimo")):
                headers[k] = v

        return HTTPRequest(
            url=url_dict,
            base_url=base_url_dict,
            headers=headers,
            query_params=query_params,
            path_params=request.path_params,
            cookies=request.cookies,
            user=_user_to_dict(request.get("user")),
            meta=_meta_to_dict(request.get("meta")),
            # Left out for now. This may contain information that the app author
            # does not want to expose.
            # session=request.session if "session" in request else {},
            # auth=request.auth if "auth" in request else {},
        )


class DebugCellCommand(Command):
    """Enter debugger mode for a cell.

    Starts the Python debugger (pdb) for the specified cell.

    Attributes:
        cell_id: Cell to debug.
        request: HTTP request context if available.
    """

    cell_id: CellId_t
    # incoming request, e.g. from Starlette or FastAPI
    request: HTTPRequest | None = None

    def __repr__(self) -> str:
        return f"DebugCellCommand(cell={self.cell_id})"


class SetBreakpointsCommand(Command):
    """Set the live debugger's breakpoints (session-scoped, not persisted).

    Replaces the full breakpoint set: the frontend always sends the complete
    map of cell id -> 1-based line numbers. Only meaningful when the
    `debugger` experimental feature is enabled.

    Attributes:
        breakpoints: Map of cell id to lines that have a breakpoint.
        request: HTTP request context if available.
    """

    breakpoints: dict[CellId_t, list[int]]
    # incoming request, e.g. from Starlette or FastAPI
    request: HTTPRequest | None = None

    def __repr__(self) -> str:
        count = sum(len(lines) for lines in self.breakpoints.values())
        return f"SetBreakpointsCommand(count={count})"


class ExecuteCellCommand(Command):
    """Execute a single cell.

    Executes a cell with the provided code. Dependent cells may be
    re-executed based on the reactive execution mode.

    Attributes:
        cell_id: Cell to execute.
        code: Python code to execute.
        request: HTTP request context if available.
        timestamp: Unix timestamp when command was created.
    """

    cell_id: CellId_t
    code: str
    # incoming request, e.g. from Starlette or FastAPI
    request: HTTPRequest | None = None
    timestamp: float = msgspec.field(default_factory=time.time)

    def __repr__(self) -> str:
        preview = self.code[:10].replace("\n", " ")
        return (
            f"ExecuteCellCommand(cell={self.cell_id}, code_preview={preview})"
        )


class ExecuteStaleCellsCommand(Command):
    """Execute all stale cells.

    Cells become stale when their dependencies change but haven't been
    re-executed yet. Brings the notebook to a consistent state.

    Attributes:
        request: HTTP request context if available.
    """

    request: HTTPRequest | None = None


class ExecuteCellsCommand(Command):
    """Execute multiple cells in a batch.

    Executes multiple cells with their corresponding code. The kernel manages
    dependency tracking and reactive execution.

    Attributes:
        cell_ids: Cells to execute.
        codes: Python code for each cell. Must match length of cell_ids.
        request: HTTP request context if available.
        timestamp: Unix timestamp when command was created.
    """

    # ids of cells to run
    cell_ids: list[CellId_t]
    # code to register/run for each cell
    codes: list[str]
    # incoming request, e.g. from Starlette or FastAPI
    request: HTTPRequest | None = None
    # time at which the request was received
    timestamp: float = msgspec.field(default_factory=time.time)

    def __repr__(self) -> str:
        return f"ExecuteCellsCommand(cells={len(self.cell_ids)})"

    @property
    def execution_requests(self) -> list[ExecuteCellCommand]:
        """Convert to individual cell execution commands.

        Returns:
            List of ExecuteCellCommand instances, one per cell.
        """
        return [
            ExecuteCellCommand(
                cell_id=cell_id,
                code=code,
                request=self.request,
                timestamp=self.timestamp,
            )
            for cell_id, code in zip(self.cell_ids, self.codes, strict=False)
        ]

    def __post_init__(self) -> None:
        assert len(self.cell_ids) == len(self.codes), (
            "Mismatched cell_ids and codes"
        )


class SyncGraphCommand(Command):
    """Synchronize the kernel graph with file manager state.

    Used when the notebook file changes externally (e.g., file reload or version control).
    Updates changed cells, deletes removed cells, and optionally executes modified cells.

    Attributes:
        cells: All cells known to file manager, mapping cell_id to code.
        run_ids: Cells to execute or update.
        delete_ids: Cells to delete from the graph.
        timestamp: Unix timestamp when command was created.
    """

    # ids of cells known to filemanager
    cells: dict[CellId_t, str]
    # From the list of ALL cells that filemanager knows about,
    # denote what should be run/ updated or deleted.
    run_ids: list[CellId_t]
    delete_ids: list[CellId_t]
    # time at which the request was received
    timestamp: float = msgspec.field(default_factory=time.time)

    @property
    def execution_requests(self) -> list[ExecuteCellCommand]:
        """Convert run_ids to individual cell execution commands.

        Returns:
            List of ExecuteCellCommand instances for cells to execute.
        """
        return [
            ExecuteCellCommand(
                cell_id=cell_id,
                code=self.cells[cell_id],
                request=None,
                timestamp=self.timestamp,
            )
            for cell_id in self.run_ids
        ]


class ExecuteScratchpadCommand(Command):
    """Execute code in the scratchpad.

    The scratchpad is a temporary execution environment that doesn't affect
    the notebook's cells or dependencies. Runs in an isolated cell with a copy
    of the global namespace, useful for experimentation.

    Attributes:
        code: Python code to execute.
        request: HTTP request context if available.
        notebook_cells: Snapshot of notebook cells from the session document.
            Used to populate the document ContextVar so code_mode can read
            cell ordering, code, names, and configs.
        cell_outputs: Snapshot of per-cell outputs (main + console) from the
            session view. Populates a parallel ContextVar so code_mode can
            expose `cell.output` and `cell.console_outputs`. Frozen at
            scratchpad start — not refreshed when `ctx.run_cell` produces
            new outputs in the same batch.
        run_id: Optional correlation ID. When set, the
            `CompletedRunNotification` emitted at the end of this command
            carries the same `run_id` so a caller holding a
            `ScratchCellListener` can filter for *its* completion and
            ignore `CompletedRun` events from unrelated commands on the
            same session.
    """

    code: str
    # incoming request, e.g. from Starlette or FastAPI
    request: HTTPRequest | None = None
    # Document snapshot — set by the execution endpoint from session.document.
    notebook_cells: tuple[NotebookCell, ...] | None = None
    # Output snapshot — set by the execution endpoint from session.session_view.
    cell_outputs: CellOutputs | None = None
    run_id: str | None = None


class RenameNotebookCommand(Command):
    """Rename or move the notebook file.

    Updates the notebook's filename in the kernel metadata.

    Attributes:
        filename: New filename or path for the notebook.
    """

    filename: str


class UpdateUIElementCommand(Command):
    """Update UI element values.

    Triggered when users interact with UI elements (sliders, inputs, dropdowns, etc.).
    Updates element values and re-executes dependent cells.

    Attributes:
        object_ids: UI elements to update.
        values: New values for the elements. Must match length of object_ids.
        request: HTTP request context if available.
        token: Unique request identifier for deduplication.
    """

    object_ids: list[UIElementId]
    values: list[Any]
    # Incoming request, e.g. from Starlette or FastAPI
    request: HTTPRequest | None = None
    # uniquely identifies the request
    token: str = msgspec.field(default_factory=lambda: str(uuid4()))

    def __repr__(self) -> str:
        return f"UpdateUIElementCommand(n_elements={len(self.object_ids)}, token={self.token})"

    def __post_init__(self) -> None:
        assert len(self.object_ids) == len(self.values), (
            "Mismatched object_ids and values"
        )
        # Empty token is not valid (but let's not fail)
        if not self.token:
            LOGGER.warning(
                "UpdateUIElementRequest with empty token is invalid"
            )

    @staticmethod
    def from_ids_and_values(
        ids_and_values: list[tuple[UIElementId, Any]],
        request: HTTPRequest | None = None,
    ) -> UpdateUIElementCommand:
        """Create command from list of (id, value) tuples.

        Args:
            ids_and_values: UI element IDs and their new values.
            request: HTTP request context if available.

        Returns:
            UpdateUIElementCommand instance.
        """
        if not ids_and_values:
            return UpdateUIElementCommand(
                object_ids=[], values=[], request=request
            )
        object_ids, values = zip(*ids_and_values, strict=False)
        return UpdateUIElementCommand(
            object_ids=list(object_ids),
            values=list(values),
            request=request,
        )

    @property
    def ids_and_values(self) -> list[tuple[UIElementId, Any]]:
        """UI element IDs and values as tuples.

        Returns:
            List of (id, value) tuples.
        """
        return list(zip(self.object_ids, self.values, strict=False))


class InvokeFunctionCommand(Command):
    """Invoke a function from a UI element.

    Called when a UI element needs to invoke a Python function.

    Attributes:
        function_call_id: Unique identifier for this call.
        namespace: Namespace where the function is registered.
        function_name: Function to invoke.
        args: Keyword arguments for the function.
    """

    function_call_id: RequestId
    namespace: str
    function_name: str
    args: dict[str, Any]

    def __repr__(self) -> str:
        return f"InvokeFunctionCommand(id={self.function_call_id}, fn={self.namespace}.{self.function_name})"


class AppMetadata(msgspec.Struct, rename="camel"):
    """Application runtime metadata.

    Attributes:
        query_params: Query parameters from the URL when running as a web app.
        cli_args: Command-line arguments passed when starting the app.
        app_config: Application-level configuration.
        argv: Full argument vector if available.
        filename: Path to the notebook file.
        docstring: Module docstring extracted from the notebook header.
    """

    query_params: SerializedQueryParams
    cli_args: SerializedCLIArgs
    app_config: _AppConfig
    argv: list[str] | None = None

    filename: str | None = None
    docstring: str | None = None


class UpdateCellConfigCommand(Command):
    """Update cell configuration.

    Updates cell-level settings like disabled state, hide code, etc.

    Attributes:
        configs: Cell IDs mapped to their config updates. Each config dict
                 can contain partial updates.
    """

    # Map from Cell ID to (possibly partial) CellConfig
    configs: dict[CellId_t, dict[str, Any]]


class UpdateUserConfigCommand(Command):
    """Update user configuration.

    Updates global marimo configuration (runtime settings, display options, editor preferences).

    Attributes:
        config: Complete user configuration.
    """

    # MarimoConfig TypedDict
    config: MarimoConfig

    def __repr__(self) -> str:
        return "UpdateUserConfigCommand(config=...)"


class CreateNotebookCommand(Command):
    """Instantiate and initialize a notebook.

    Sent when a notebook is first loaded. Contains all cells and initial UI element values.

    Attributes:
        execution_requests: ExecuteCellCommand for each notebook cell.
        cell_ids: Initial cell IDs in the notebook.
        set_ui_element_value_request: Initial UI element values.
        auto_run: Whether to automatically execute cells on instantiation.
        request: HTTP request context if available.
    """

    execution_requests: tuple[ExecuteCellCommand, ...]
    cell_ids: tuple[CellId_t, ...]
    set_ui_element_value_request: UpdateUIElementCommand
    auto_run: bool
    request: HTTPRequest | None = None


class DeleteCellCommand(Command):
    """Delete a cell from the notebook.

    Removes cell from the dependency graph and cleans up its variables.
    Dependent cells may become stale.

    Attributes:
        cell_id: Cell to delete.
    """

    cell_id: CellId_t


class StopKernelCommand(Command):
    """Stop kernel execution.

    Signals the kernel to stop processing and shut down gracefully.
    Used when closing a notebook or terminating a session.
    """


class CodeCompletionCommand(Command):
    """Request code completion suggestions.

    Sent when the user requests autocomplete. Provides code context up to
    the cursor position for the language server.

    Attributes:
        id: Unique identifier for this request.
        document: Source code up to the cursor position.
        cell_id: Cell where completion is requested.
    """

    id: RequestId
    document: str
    """Source code found in the cell up to the cursor position."""
    cell_id: CellId_t

    def __repr__(self) -> str:
        return f"CodeCompletionCommand(id={self.id}, cell={self.cell_id})"


class InstallPackagesCommand(Command):
    """Install Python packages.

    Installs missing packages using the specified package manager. Triggered
    automatically on import errors or manually by the user.

    Attributes:
        manager: Package manager to use ('pip', 'conda', 'uv', etc.).
        versions: Package names mapped to version specifiers. Empty version
                  means install latest.
        source: Where to install. "kernel" (default) dispatches to the kernel
                subprocess; "server" installs directly into the server's Python
                environment (sys.executable), used when the server itself needs
                a package (e.g. nbformat for IPYNB auto-export in sandbox mode).
    """

    # TODO: index URL (index/channel/...)
    manager: str

    # Map from package name to desired version
    # If the package name is not in the map, the latest version
    # will be installed
    versions: dict[str, str]

    source: Literal["kernel", "server"] = "kernel"


class PreviewDatasetColumnCommand(Command):
    """Preview a dataset column.

    Retrieves and displays data from a single column (dataframe or SQL table).
    Used by the data explorer UI.

    Attributes:
        source_type: Data source type ('dataframe', 'sql', etc.).
        source: Source identifier (connection string or variable name).
        table_name: Table or dataframe variable name.
        column_name: Column to preview.
        fully_qualified_table_name: Full database.schema.table name for SQL.
    """

    # The source type of the dataset
    source_type: DataTableSource
    # The source of the dataset
    source: str
    # The name of the dataset
    # If this is a Python dataframe, this is the variable name
    # If this is an SQL table, this is the table name
    table_name: str
    # The name of the column
    column_name: str
    # The fully qualified name of the table
    # This is the database.schema.table name
    fully_qualified_table_name: str | None = None


class PreviewSQLTableCommand(Command):
    """Preview SQL table details.

    Retrieves metadata and sample data for a table. Used by the SQL editor
    and data explorer.

    Attributes:
        request_id: Unique identifier for this request.
        engine: SQL engine ('postgresql', 'mysql', 'duckdb', etc.).
        database: Database containing the table.
        schema: Schema containing the table.
        table_name: Table to preview.
        schema_path: Path of nested schemas (relative to `database`) for
            catalogs with nested schemas. Empty for the top level.
    """

    request_id: RequestId
    engine: str
    database: str
    schema: str
    table_name: str
    schema_path: list[str] = msgspec.field(default_factory=list)


class ListSQLTablesCommand(Command):
    """List tables in an SQL schema.

    Retrieves names of all tables and views in a schema. Used by the SQL
    editor for table selection.

    Attributes:
        request_id: Unique identifier for this request.
        engine: SQL engine ('postgresql', 'mysql', 'duckdb', etc.).
        database: Database to query.
        schema: Schema to list tables from.
        schema_path: Path of nested schemas (relative to `database`) for
            catalogs with nested schemas. Empty for the top level.
    """

    request_id: RequestId
    engine: str
    database: str
    schema: str
    schema_path: list[str] = msgspec.field(default_factory=list)


class ListSQLSchemasCommand(Command):
    """List schemas in an SQL database.

    Retrieves names of all schemas in a database. Used by the SQL editor for
    schema selection.

    Attributes:
        request_id: Unique identifier for this request.
        engine: SQL engine ('postgresql', 'mysql', 'duckdb', etc.).
        database: Database to query.
        schema_path: Parent schema path whose child schemas to list.
            Empty lists the database's top-level schemas.
    """

    request_id: RequestId
    engine: str
    database: str
    schema_path: list[str] = msgspec.field(default_factory=list)


class ListDataSourceConnectionCommand(Command):
    """List data source schemas.

    Retrieves available schemas for a data source engine.

    Attributes:
        engine: Data source engine identifier.
    """

    engine: str


class ValidateSQLCommand(Command):
    """Validate an SQL query.

    Checks if an SQL query is valid by parsing against a dialect (no DB connection)
    or validating against an actual database.

    Attributes:
        request_id: Unique identifier for this request.
        query: SQL query to validate.
        only_parse: If True, only parse using dialect. If False, validate against DB.
        engine: SQL engine (required if only_parse is False).
        dialect: SQL dialect for parsing (required if only_parse is True).
    """

    request_id: RequestId
    query: str
    # Whether to only parse the query or validate against the database
    # Parsing is done without a DB connection and uses dialect, whereas validation requires a connection
    only_parse: bool
    engine: str | None = None
    dialect: str | None = None


class StorageListEntriesCommand(Command):
    """List storage entries at a prefix.

    Navigates storage like a folder tree using delimiter-based listing.
    Returns entries (files/objects) and virtual directories at one level.

    Attributes:
        request_id: Unique identifier for this request.
        namespace: Variable name identifying the storage backend.
        limit: Max entries to return.
        prefix: Path prefix to list (None = root).
        page_token: Token for the next page of entries.
    """

    request_id: RequestId
    namespace: str
    limit: int
    prefix: str | None = None
    page_token: str | None = None


class StorageDownloadCommand(Command):
    """Download a storage entry.

    Obtains a pre-signed URL or downloads the file locally and returns a virtual file URL
    so the frontend can fetch the contents.

    Attributes:
        request_id: Unique identifier for this request.
        namespace: Variable name identifying the storage backend.
        path: Full path of the entry to download.
        preview: If true, a local preview of the file is returned.
            This is useful if you need to bypass CORS.
    """

    request_id: RequestId
    namespace: str
    path: str
    preview: bool = False


class DiscoverDataSourcesCommand(Command):
    """Discover datasource connections from the live kernel environment and configuration.

    Attributes:
        request_id: Unique identifier for this request.
    """

    request_id: RequestId


class ListSecretKeysCommand(Command):
    """List available secret keys.

    Retrieves secret names without exposing values.

    Attributes:
        request_id: Unique identifier for this request.
    """

    request_id: RequestId


class ModelUpdateMessage(
    msgspec.Struct, tag="update", tag_field="method", rename="camel"
):
    """Widget model state update message.

    Attributes:
        state: Model state updates.
        buffer_paths: Paths within state dict pointing to binary buffers.
    """

    state: dict[str, Encodable]
    buffer_paths: list[list[str | int]]

    def into_comm_payload_content(self) -> dict[str, Any]:
        return {
            "data": {
                "method": "update",
                "state": self.state,
                "buffer_paths": self.buffer_paths,
            }
        }


class ModelCustomMessage(
    msgspec.Struct, tag="custom", tag_field="method", rename="camel"
):
    """Custom widget message.

    Attributes:
        content: Arbitrary content for the custom message.
    """

    content: Encodable

    def into_comm_payload_content(self) -> dict[str, Any]:
        return {
            "data": {
                "method": "custom",
                "content": self.content,
            }
        }


ModelMessage = ModelUpdateMessage | ModelCustomMessage


class ModelCommand(Command):
    """Widget model message command.

    Handles widget model communication between frontend and backend.

    Attributes:
        model_id: Widget model identifier.
        message: Model message (update or custom).
        buffers: Base64-encoded binary buffers.
        token: Unique identifier for deduplication across dual queues.
    """

    model_id: WidgetModelId
    message: ModelMessage
    buffers: list[bytes]
    token: str = msgspec.field(default_factory=lambda: str(uuid4()))

    def into_comm_payload(self) -> dict[str, Any]:
        return {
            "content": self.message.into_comm_payload_content(),
            "buffers": self.buffers,
        }


# Commands that can be batched and merged (last-write-wins) by the
# SetUIElementRequestManager to avoid redundant cell re-executions.
BatchableCommand = UpdateUIElementCommand | ModelCommand


class RefreshSecretsCommand(Command):
    """Refresh secrets from the secrets store.

    Reloads secrets from the provider without restarting the kernel.
    """


class ClearCacheCommand(Command):
    """Clear all cached data.

    Clears all cache contexts, freeing memory and disk space.
    Affects all cells using the @cache decorator.
    """


class GetCacheInfoCommand(Command):
    """Retrieve cache statistics.

    Collects cache usage info across all contexts (hit/miss rates, time saved, disk usage).
    """


CommandMessage = (
    # Notebook operations
    CreateNotebookCommand
    | RenameNotebookCommand
    | CodeCompletionCommand
    # Cell execution and management
    | ExecuteCellsCommand
    | ExecuteScratchpadCommand
    | ExecuteStaleCellsCommand
    | DebugCellCommand
    | SetBreakpointsCommand
    | DeleteCellCommand
    | SyncGraphCommand
    | UpdateCellConfigCommand
    # Package management
    | InstallPackagesCommand
    # UI element and widget model operations
    | UpdateUIElementCommand
    | ModelCommand
    | InvokeFunctionCommand
    # User/configuration operations
    | UpdateUserConfigCommand
    # Data SQL operations
    | PreviewDatasetColumnCommand
    | PreviewSQLTableCommand
    | ListSQLTablesCommand
    | ListSQLSchemasCommand
    | ValidateSQLCommand
    | ListDataSourceConnectionCommand
    | DiscoverDataSourcesCommand
    # Storage operations
    | StorageListEntriesCommand
    | StorageDownloadCommand
    # Secrets management
    | ListSecretKeysCommand
    | RefreshSecretsCommand
    # Cache management
    | ClearCacheCommand
    | GetCacheInfoCommand
    # Kernel operations
    | StopKernelCommand
)
"""Union of all command messages.

All commands that can be sent to the kernel.

"""


OutOfBandCommand = CodeCompletionCommand | SetBreakpointsCommand
"""Commands processed off the main control loop.

Unlike the rest of `CommandMessage` (which the kernel handles serially, and so
cannot be delivered while a cell is executing), these are drained by a
background worker and applied immediately. A command belongs here when it is
fire-and-forget, cheap to apply, and useful mid-execution. Add a member to this
union and a branch to `Kernel.dispatch_out_of_band` to introduce a new one.
"""
