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

from typing import Any, Literal

import msgspec

from marimo._ast.cell import CellConfig
from marimo._messaging.notebook.changes import DocumentChange
from marimo._runtime.commands import (
    ClearCacheCommand,
    CodeCompletionCommand,
    DebugCellCommand,
    DeleteCellCommand,
    DiscoverDataSourcesCommand,
    ExecuteCellsCommand,
    ExecuteScratchpadCommand,
    GetCacheInfoCommand,
    HTTPRequest,
    InstallPackagesCommand,
    InvokeFunctionCommand,
    ListDataSourceConnectionCommand,
    ListSecretKeysCommand,
    ListSQLSchemasCommand,
    ListSQLTablesCommand,
    ModelCommand,
    PreviewDatasetColumnCommand,
    PreviewSQLTableCommand,
    SetBreakpointsCommand,
    StorageDownloadCommand,
    StorageListEntriesCommand,
    UpdateCellConfigCommand,
    UpdateUIElementCommand,
    UpdateUserConfigCommand,
    ValidateSQLCommand,
)
from marimo._types.ids import CellId_t


class ListSecretKeysRequest(ListSecretKeysCommand, tag=False):
    def as_command(self) -> ListSecretKeysCommand:
        return ListSecretKeysCommand(request_id=self.request_id)


class DiscoverDataSourcesRequest(DiscoverDataSourcesCommand, tag=False):
    def as_command(self) -> DiscoverDataSourcesCommand:
        return DiscoverDataSourcesCommand(request_id=self.request_id)


class ClearCacheRequest(ClearCacheCommand, tag=False):
    def as_command(self) -> ClearCacheCommand:
        return ClearCacheCommand()


class GetCacheInfoRequest(GetCacheInfoCommand, tag=False):
    def as_command(self) -> GetCacheInfoCommand:
        return GetCacheInfoCommand()


class DebugCellRequest(DebugCellCommand, tag=False):
    def as_command(self) -> DebugCellCommand:
        return DebugCellCommand(cell_id=self.cell_id, request=self.request)


class SetBreakpointsRequest(SetBreakpointsCommand, tag=False):
    def as_command(self) -> SetBreakpointsCommand:
        return SetBreakpointsCommand(
            breakpoints=self.breakpoints, request=self.request
        )


class ExecuteScratchpadRequest(ExecuteScratchpadCommand, tag=False):
    def as_command(self) -> ExecuteScratchpadCommand:
        return ExecuteScratchpadCommand(code=self.code, request=self.request)


class InvokeFunctionRequest(InvokeFunctionCommand, tag=False):
    def as_command(self) -> InvokeFunctionCommand:
        return InvokeFunctionCommand(
            function_call_id=self.function_call_id,
            namespace=self.namespace,
            function_name=self.function_name,
            args=self.args,
        )


class UpdateUIElementRequest(UpdateUIElementCommand, tag=False):
    def as_command(self) -> UpdateUIElementCommand:
        return UpdateUIElementCommand(
            object_ids=self.object_ids,
            values=self.values,
            request=self.request,
            token=self.token,
        )


class ModelRequest(ModelCommand, tag=False):
    def as_command(self) -> ModelCommand:
        return ModelCommand(
            model_id=self.model_id,
            message=self.message,
            buffers=self.buffers,
        )


class ListDataSourceConnectionRequest(
    ListDataSourceConnectionCommand, tag=False
):
    def as_command(self) -> ListDataSourceConnectionCommand:
        return ListDataSourceConnectionCommand(engine=self.engine)


class ListSQLTablesRequest(ListSQLTablesCommand, tag=False):
    def as_command(self) -> ListSQLTablesCommand:
        return ListSQLTablesCommand(
            request_id=self.request_id,
            engine=self.engine,
            database=self.database,
            schema=self.schema,
            schema_path=self.schema_path,
        )


class ListSQLSchemasRequest(ListSQLSchemasCommand, tag=False):
    def as_command(self) -> ListSQLSchemasCommand:
        return ListSQLSchemasCommand(
            request_id=self.request_id,
            engine=self.engine,
            database=self.database,
            schema_path=self.schema_path,
        )


class PreviewDatasetColumnRequest(PreviewDatasetColumnCommand, tag=False):
    def as_command(self) -> PreviewDatasetColumnCommand:
        return PreviewDatasetColumnCommand(
            source_type=self.source_type,
            source=self.source,
            table_name=self.table_name,
            column_name=self.column_name,
            fully_qualified_table_name=self.fully_qualified_table_name,
        )


class PreviewSQLTableRequest(PreviewSQLTableCommand, tag=False):
    def as_command(self) -> PreviewSQLTableCommand:
        return PreviewSQLTableCommand(
            request_id=self.request_id,
            engine=self.engine,
            database=self.database,
            schema=self.schema,
            table_name=self.table_name,
            schema_path=self.schema_path,
        )


class ValidateSQLRequest(ValidateSQLCommand, tag=False):
    def as_command(self) -> ValidateSQLCommand:
        return ValidateSQLCommand(
            query=self.query,
            only_parse=self.only_parse,
            engine=self.engine,
            dialect=self.dialect,
            request_id=self.request_id,
        )


class StorageListEntriesRequest(StorageListEntriesCommand, tag=False):
    def as_command(self) -> StorageListEntriesCommand:
        return StorageListEntriesCommand(
            request_id=self.request_id,
            namespace=self.namespace,
            limit=self.limit,
            prefix=self.prefix,
            page_token=self.page_token,
        )


class StorageDownloadRequest(StorageDownloadCommand, tag=False):
    def as_command(self) -> StorageDownloadCommand:
        return StorageDownloadCommand(
            request_id=self.request_id,
            namespace=self.namespace,
            path=self.path,
            preview=self.preview,
        )


class UpdateUserConfigRequest(UpdateUserConfigCommand, tag=False):
    def as_command(self) -> UpdateUserConfigCommand:
        return UpdateUserConfigCommand(config=self.config)


class DeleteCellRequest(DeleteCellCommand, tag=False):
    def as_command(self) -> DeleteCellCommand:
        return DeleteCellCommand(cell_id=self.cell_id)


class InstallPackagesRequest(InstallPackagesCommand, tag=False):
    def as_command(self) -> InstallPackagesCommand:
        return InstallPackagesCommand(
            manager=self.manager, versions=self.versions, source=self.source
        )


class UpdateCellConfigRequest(UpdateCellConfigCommand, tag=False):
    def as_command(self) -> UpdateCellConfigCommand:
        return UpdateCellConfigCommand(configs=self.configs)


class CodeCompletionRequest(CodeCompletionCommand, tag=False):
    def as_command(self) -> CodeCompletionCommand:
        return CodeCompletionCommand(
            id=self.id,
            document=self.document,
            cell_id=self.cell_id,
        )


class BaseResponse(msgspec.Struct, rename="camel"):
    success: bool


class SuccessResponse(BaseResponse):
    success: bool = True


class ErrorResponse(BaseResponse):
    success: bool = False
    message: str | None = None


class KernelStatusResponse(msgspec.Struct, rename="camel"):
    # `running`: at least one cell is queued or running.
    # `idle`: the kernel is alive but not executing.
    # `stopped`: the kernel process is not running (dead or not started).
    state: Literal["running", "idle", "stopped"]


class FormatCellsRequest(msgspec.Struct, rename="camel"):
    codes: dict[CellId_t, str]
    line_length: int


class FormatResponse(msgspec.Struct, rename="camel"):
    codes: dict[CellId_t, str]


class ReadCodeResponse(msgspec.Struct, rename="camel"):
    contents: str


class RenameNotebookRequest(msgspec.Struct, rename="camel"):
    filename: str


class NotebookDocumentTransactionRequest(msgspec.Struct, rename="camel"):
    changes: list[DocumentChange]


class FocusCellRequest(msgspec.Struct, rename="camel"):
    cell_id: CellId_t


class ExecuteCellsRequest(msgspec.Struct, rename="camel"):
    # 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

    def as_command(self) -> ExecuteCellsCommand:
        return ExecuteCellsCommand(
            cell_ids=self.cell_ids,
            codes=self.codes,
            request=self.request,
        )

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


class SaveNotebookRequest(msgspec.Struct, rename="camel"):
    # id of each cell
    cell_ids: list[CellId_t]
    # code for each cell
    codes: list[str]
    # name of each cell
    names: list[str]
    # config for each cell
    configs: list[CellConfig]
    # path to app
    filename: str
    # layout of app
    layout: dict[str, Any] | None = None
    # persist the file to disk
    persist: bool = True

    # Validate same length
    def __post_init__(self) -> None:
        assert len(self.cell_ids) == len(self.codes), (
            "Mismatched cell_ids and codes"
        )
        assert len(self.cell_ids) == len(self.names), (
            "Mismatched cell_ids and names"
        )
        assert len(self.cell_ids) == len(self.configs), (
            "Mismatched cell_ids and configs"
        )


class CopyNotebookRequest(msgspec.Struct, rename="camel"):
    # path to app
    source: str
    destination: str


class SaveAppConfigurationRequest(msgspec.Struct, rename="camel"):
    # partial app config
    config: dict[str, Any]


class SaveUserConfigurationRequest(msgspec.Struct, rename="camel"):
    # deep partial user configuration; keys with value `None` are removed
    # from the on-disk merged config (None-as-delete)
    config: dict[str, Any]


class StdinRequest(msgspec.Struct, rename="camel"):
    text: str


class InvokeAiToolRequest(msgspec.Struct, rename="camel"):
    tool_name: str
    arguments: dict[str, Any]


class InvokeAiToolResponse(BaseResponse):
    tool_name: str
    result: Any
    error: str | None = None


class MCPStatusResponse(msgspec.Struct, rename="camel"):
    status: Literal["ok", "partial", "error"]
    error: str | None = None
    servers: dict[
        str, Literal["pending", "connected", "disconnected", "failed"]
    ] = {}  # server_name -> status


class MCPRefreshResponse(BaseResponse):
    error: str | None = None
    servers: dict[str, bool] = {}  # server_name -> connected
