from __future__ import annotations

import uuid
from collections.abc import Callable, Sequence
from typing import Any, Generic, Literal, Self, TypeVar, cast, overload

from pydantic import BaseModel, ConfigDict, Field, model_validator

from mistralai.workflows.conversational_ui_components import LLM_UI_VERSION, UIComponent
from mistralai.workflows.core.task.task import Task

ChoiceT = TypeVar("ChoiceT", bound=str)


class TextChunk(BaseModel):
    type: Literal["text"] = "text"
    text: str


class SearchReplaceBlock(BaseModel):
    search: str
    replace: str


class CreateFileOperation(BaseModel):
    type: Literal["create"] = "create"
    uri: str
    content: str


class ReplaceFileOperation(BaseModel):
    type: Literal["replace"] = "replace"
    uri: str
    fileContentBefore: str
    blocks: list[SearchReplaceBlock]


class DeleteFileOperation(BaseModel):
    type: Literal["delete"] = "delete"
    uri: str


FileOperation = CreateFileOperation | ReplaceFileOperation | DeleteFileOperation


class FileToolUIState(BaseModel):
    type: Literal["file"] = "file"
    toolCallId: str
    operations: list[FileOperation]


class ToolResultPending(BaseModel):
    status: Literal["pending"] = "pending"


class ToolResultRunning(BaseModel):
    status: Literal["running"] = "running"


class ToolResultFailed(BaseModel):
    status: Literal["failed"] = "failed"
    error: Any


class ToolResultSuccess(BaseModel):
    status: Literal["success"] = "success"
    value: Any


ToolResult = ToolResultPending | ToolResultRunning | ToolResultFailed | ToolResultSuccess


class GenericToolUIState(BaseModel):
    type: Literal["generic_tool"] = "generic_tool"
    toolCallId: str
    name: str
    arguments: dict[str, Any]
    result: ToolResult


class CommandResultRunning(BaseModel):
    status: Literal["running"] = "running"


class CommandResultSuccess(BaseModel):
    status: Literal["success"] = "success"
    output: str


class CommandResultFailed(BaseModel):
    status: Literal["failed"] = "failed"
    error: str


CommandResult = CommandResultRunning | CommandResultSuccess | CommandResultFailed


class CommandToolUIState(BaseModel):
    type: Literal["command"] = "command"
    toolCallId: str
    command: str
    result: CommandResult


ToolUIState = FileToolUIState | GenericToolUIState | CommandToolUIState


class ChatAssistantWorkingTask(BaseModel):
    model_config = ConfigDict(title="working")

    type: Literal["tool", "thinking"] | str = "tool"
    title: str
    content: str
    toolUIState: ToolUIState | None = None


class TodoListItemState(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    status: TodoListItemStatus = "todo"
    title: str
    description: str


class TodoListState(BaseModel):
    model_config = ConfigDict(title="todo_list")
    items: list[TodoListItemState]


type TodoListItemStatus = Literal["todo", "in_progress", "done"]


class TodoListItem:
    """A todo list item with context manager support for automatic status updates.

    Usage:
        ```python
        item = TodoListItem(title="Validate expense", description="Check details")

        async with TodoList(items=[item]) as todo_list:
            # Using context manager for automatic status transitions
            async with item:
                # Status is set to "in_progress" on enter
                ...
                # Status is set to "done" on successful exit

            # Or manual status control
            await item.set_status("in_progress")
            ...
            await item.set_status("done")
        ```
    """

    def __init__(self, title: str, description: str) -> None:
        self._id: str = str(uuid.uuid4())
        self._title: str = title
        self._description: str = description
        self._status: TodoListItemStatus = "todo"
        self._todo_list: TodoList | None = None

    @property
    def id(self) -> str:
        return self._id

    @property
    def title(self) -> str:
        return self._title

    @property
    def description(self) -> str:
        return self._description

    @property
    def status(self) -> TodoListItemStatus:
        return self._status

    def _bind(self, todo_list: TodoList) -> None:
        """Internal: bind this item to a TodoList."""
        if self._todo_list is not None and self._todo_list is not todo_list:
            raise RuntimeError("TodoListItem is already bound to a different TodoList")
        self._todo_list = todo_list

    def _unbind(self) -> None:
        """Internal: unbind this item from a TodoList."""
        self._todo_list = None

    async def set_status(self, status: TodoListItemStatus) -> None:
        """Update the status of this item and emit an event."""
        if self._todo_list is None:
            raise RuntimeError("TodoListItem is not bound to a TodoList - use within a TodoList context")
        self._status = status
        await self._todo_list._sync_state()

    async def __aenter__(self) -> TodoListItem:
        """Set status to 'in_progress' on enter."""
        await self.set_status("in_progress")
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: Any | None,
    ) -> None:
        """Set status to 'done' on success, leave unchanged on exception."""
        if exc_type is None:
            await self.set_status("done")


class TextOutput(BaseModel):
    type: Literal["text"] = "text"
    text: str


class CanvasPayload(BaseModel):
    type: Literal[
        "text/markdown",
        "text/html",
        "image/svg+xml",
        "slides",
        "react",
        "code",
        "mermaid",
    ]
    title: str
    content: str
    language: str | None = None


class CanvasResource(BaseModel):
    uri: str = Field(default_factory=lambda: f"file://canvas/{uuid.uuid4()}")
    mimeType: Literal["application/vnd.mistral.canvas"] = "application/vnd.mistral.canvas"
    canvas: CanvasPayload


class UIComponentResource(BaseModel):
    mimeType: Literal["application/vnd.mistral.ui-component"] = "application/vnd.mistral.ui-component"
    component: UIComponent
    version: str = LLM_UI_VERSION
    display: Literal["inline", "block"] = "block"


class ResourceOutput(BaseModel):
    type: Literal["resource"] = "resource"
    resource: CanvasResource | UIComponentResource


ContentChunk = TextOutput | ResourceOutput


class _AssistantMessageState(BaseModel):
    """Internal state for assistant message custom task events."""

    model_config = ConfigDict(title="assistant_message")
    contentChunks: list[ContentChunk] = Field(default_factory=list)


@overload
async def send_assistant_message(
    text: str,
    *,
    canvas: CanvasResource | None = None,
) -> None: ...


@overload
async def send_assistant_message(
    text: Sequence[ContentChunk],
) -> None: ...


async def send_assistant_message(
    text: str | Sequence[ContentChunk],
    *,
    canvas: CanvasResource | None = None,
) -> None:
    """Send an assistant message to the user.

    This function allows workflows to send messages to users before requesting input,
    improving the conversational experience. The message appears in the chat as if
    sent by an assistant.

    Args:
        text: Either a text string or a sequence of content chunks.
        canvas: Optional canvas resource to include alongside the text (only when passing a string).

    Example:
        ```python
        await send_assistant_message("Please provide your favorite Pokemon:")

        await send_assistant_message(
            "Here is your document draft:",
            canvas=CanvasResource(
                canvas=CanvasPayload(type="text/markdown", title="Draft", content="# Hello"),
            ),
        )

        await send_assistant_message([
            TextOutput(text="Here is the report:"),
            ResourceOutput(resource=CanvasResource(
                canvas=CanvasPayload(type="text/markdown", title="Draft", content="# Hello"),
            )),
        ])
        ```
    """
    if isinstance(text, str):
        chunks: list[ContentChunk] = [TextOutput(text=text)]
        if canvas is not None:
            chunks.append(ResourceOutput(resource=canvas))
    else:
        chunks = list(text)
    async with Task[_AssistantMessageState](
        type="assistant_message",
        state=_AssistantMessageState(contentChunks=chunks),
    ):
        pass


class TodoList(Task[TodoListState]):
    """A list of todo items with real-time status updates.

    TodoList is a specialized Task that manages a collection of TodoListItem instances
    and emits events as items progress through their lifecycle.

    Usage:
        ```python
        item1 = TodoListItem(title="Validate", description="Check details")
        item2 = TodoListItem(title="Process", description="Execute action")

        async with TodoList(items=[item1, item2]) as todo_list:
            # Add items
            item3 = TodoListItem(title="Notify", description="Send notification")
            await todo_list.add_item(item3)

            # Remove items
            await todo_list.remove_item(item2)

            # Work through items
            async with item1:
                ...  # item1 goes in_progress -> done
        ```
    """

    def __init__(self, items: list[TodoListItem] | None = None, id: str | None = None) -> None:
        self._todo_items: list[TodoListItem] = list(items) if items else []
        for item in self._todo_items:
            item._bind(self)
        super().__init__(
            type="todo_list",
            state=self._build_state(),
            id=id,
        )

    def _build_state(self) -> TodoListState:
        return TodoListState(
            items=[
                TodoListItemState(
                    id=item.id,
                    status=item.status,
                    title=item.title,
                    description=item.description,
                )
                for item in self._todo_items
            ]
        )

    async def _sync_state(self) -> None:
        await self.set_state(self._build_state())

    @property
    def items(self) -> list[TodoListItem]:
        return list(self._todo_items)

    async def add_item(self, item: TodoListItem) -> None:
        if item in self._todo_items:
            raise ValueError("TodoListItem is already in this TodoList")
        item._bind(self)
        self._todo_items.append(item)
        await self._sync_state()

    async def remove_item(self, item: TodoListItem) -> None:
        if item not in self._todo_items:
            raise ValueError("TodoListItem is not in this TodoList")
        item._unbind()
        self._todo_items.remove(item)
        await self._sync_state()

    async def __aenter__(self) -> Self:
        await super().__aenter__()
        return self


class ChatAssistantWorkflowOutput(BaseModel):
    content: list[ContentChunk] = Field(default_factory=list)
    structuredContent: dict[str, Any] | None = None
    isError: bool | None = None


ChatInputMessageContent = list[TextChunk]

_ModelT = TypeVar("_ModelT", bound=BaseModel)


def input_tag(tag: str) -> Callable[[type[_ModelT]], type[_ModelT]]:
    """Annotate a Pydantic model's JSON schema with ``x-input-tag``.

    Clients use this annotation to select the correct input variant when a
    workflow exposes multiple input schemas for different consumers.

    Args:
        tag: Arbitrary identifier surfaced in the schema as ``x-input-tag``.

    Example::

        @input_tag("my-custom-input")
        class MyWorkflowInput(BaseModel):
            message: ChatInputMessageContent
    """

    def _decorator(cls: type[_ModelT]) -> type[_ModelT]:
        existing = cls.model_config.get("json_schema_extra")
        if callable(existing):
            original_callable = cast("Callable[[dict[str, Any]], None]", existing)

            def _composed(schema: dict[str, Any]) -> None:
                original_callable(schema)
                schema["x-input-tag"] = tag

            cls.model_config["json_schema_extra"] = _composed
        else:
            existing_dict = existing if isinstance(existing, dict) else {}
            cls.model_config["json_schema_extra"] = {**existing_dict, "x-input-tag": tag}
        cls.model_rebuild(force=True)
        return cls

    return _decorator


class ChatInputModel(BaseModel):
    model_config = ConfigDict(title="ChatInput", extra="forbid")
    message: ChatInputMessageContent


def ChatInput(
    prompt: str = "Type your message", *, suggestions: list[ChatInputMessageContent] | None = None
) -> type[ChatInputModel]:
    """Create a ChatInput schema with a custom prompt/placeholder.

    Args:
        prompt: The placeholder text shown in the chat input field.
        suggestions: List of message suggestions that can be directly submitted.

    Example:
        ```python
        user_input = await self.wait_for_input(ChatInput(
            "Ask a follow-up question",
            suggestions=[[TextChunk(text="What is the option 1?")], [TextChunk(text="What is the option 2?")]],
        ))
        ```
    """

    class _ChatInput(ChatInputModel):
        message: list[TextChunk] = Field(description=prompt, examples=suggestions)

    return _ChatInput


class CanvasInputData(BaseModel):
    title: str
    content: str


class CanvasInputModel(BaseModel):
    model_config = ConfigDict(title="Canvas", extra="forbid")
    canvas: CanvasInputData
    chatInput: ChatInputModel | None = None


def CanvasInput(canvas_uri: str, *, prompt: str | None = None) -> type[CanvasInputModel]:
    """Create a CanvasInput schema that references a previously output canvas.

    The canvas identified by ``canvas_uri`` must have been output in a prior
    workflow step with a ``CanvasResource`` whose ``uri`` matches.

    Args:
        canvas_uri: URI of the canvas to reference (must match a
            ``CanvasResource.uri`` from a previous workflow step output).
        prompt: Optional placeholder text for a chat input field alongside
            the canvas. When provided, the user can send follow-up
            instructions with their canvas edits.

    Example:
        ```python
        result = await self.wait_for_input(
            CanvasInput(canvas_uri="file://canvas/draft.md")
        )
        print(result.canvas.title, result.canvas.content)

        result = await self.wait_for_input(
            CanvasInput(
                canvas_uri="file://canvas/draft.md",
                prompt="Any feedback on the document?",
            )
        )
        if result.chatInput:
            print(result.chatInput.message)
        ```
    """
    metadata = {"canvasUri": canvas_uri}

    def _canvas_schema_extra(schema: dict[str, Any]) -> None:
        schema["$metadata"] = metadata
        props = schema.get("properties", {})
        if prompt is None:
            props.pop("chatInput", None)

    class _ChatInput(ChatInputModel):
        message: list[TextChunk] = Field(description=prompt or "")

    class _CanvasInput(CanvasInputModel):
        model_config = ConfigDict(title="Canvas", json_schema_extra=_canvas_schema_extra)
        chatInput: _ChatInput | None = None

    return _CanvasInput


class _SingleChoiceSchemaExtra:
    def __init__(self, options: Sequence[tuple[str, str]] | Sequence[str], description: str | None):
        self.allowed_values: list[str] = []
        for opt in options:
            if isinstance(opt, str):
                self.allowed_values.append(opt)
            else:
                self.allowed_values.append(opt[0])
        self.options = options
        self.description = description

    def __call__(self, schema: dict[str, Any]) -> None:
        one_of = []
        for opt in self.options:
            if isinstance(opt, str):
                one_of.append({"type": "string", "const": opt})
            else:
                value, label = opt
                one_of.append({"type": "string", "const": value, "description": label})
        schema["oneOf"] = one_of
        if self.description:
            schema["description"] = self.description
        schema.pop("type", None)
        schema.pop("title", None)


class _MultiChoiceSchemaExtra:
    def __init__(self, options: list[tuple[str, str]] | list[str], description: str | None):
        self.allowed_values: list[str] = []
        for opt in options:
            if isinstance(opt, str):
                self.allowed_values.append(opt)
            else:
                self.allowed_values.append(opt[0])
        if len(self.allowed_values) != len(set(self.allowed_values)):
            raise ValueError("MultiChoice options contain duplicate values")
        self.options = options
        self.description = description

    def __call__(self, schema: dict[str, Any]) -> None:
        any_of = []
        for opt in self.options:
            if isinstance(opt, str):
                any_of.append({"type": "string", "const": opt})
            else:
                value, label = opt
                any_of.append({"type": "string", "const": value, "description": label})
        schema["type"] = "array"
        schema["items"] = {"anyOf": any_of}
        schema["minItems"] = 1
        schema["uniqueItems"] = True
        if self.description:
            schema["description"] = self.description
        schema.pop("title", None)


class FormInput(BaseModel):
    """Base class for form input models used with wait_for_input().

    FormInput models generate JSON schemas that are compatible with
    the workflow form UI. Use this as a base class for models that
    will be used with wait_for_input().

    Example:
        ```python
        class ExpenseForm(FormInput):
            reason: str = TextField(description="Expense reason")
            amount: float = NumberField(description="Amount in USD", minimum=0)
            category: str = SingleChoice(
                options=[("travel", "Travel"), ("equipment", "Equipment")],
                description="Category"
            )
        ```
    """

    model_config = ConfigDict(title="FormInput")

    @model_validator(mode="after")
    def _validate_single_choice_fields(self) -> Self:
        for field_name, field_info in self.__class__.model_fields.items():
            json_schema_extra = field_info.json_schema_extra
            if isinstance(json_schema_extra, _SingleChoiceSchemaExtra):
                value = getattr(self, field_name)
                if value not in json_schema_extra.allowed_values:
                    raise ValueError(
                        f"Invalid value '{value}' for field '{field_name}'. "
                        f"Must be one of: {json_schema_extra.allowed_values}"
                    )
        return self

    @model_validator(mode="after")
    def _validate_multi_choice_fields(self) -> Self:
        for field_name, field_info in self.__class__.model_fields.items():
            json_schema_extra = field_info.json_schema_extra
            if isinstance(json_schema_extra, _MultiChoiceSchemaExtra):
                values = getattr(self, field_name)
                if not isinstance(values, list):
                    raise ValueError(f"Invalid value for field '{field_name}'. Expected a list.")
                if len(values) == 0:
                    raise ValueError(f"Invalid value for field '{field_name}'. List must not be empty.")
                if len(values) != len(set(values)):
                    raise ValueError(f"Invalid value for field '{field_name}'. List must not contain duplicates.")
                for v in values:
                    if v not in json_schema_extra.allowed_values:
                        raise ValueError(
                            f"Invalid value '{v}' for field '{field_name}'. "
                            f"Must be one of: {json_schema_extra.allowed_values}"
                        )
        return self


def TextField(
    description: str,
    *,
    pattern: str | None = None,
) -> Any:
    """Create a text field for form inputs.

    Args:
        description: Human-readable description shown as label/placeholder.
        pattern: Optional regex pattern for validation.

    Example:
        ```python
        class MyForm(FormInput):
            name: str = TextField(description="Enter your name")
        ```
    """

    def _schema_extra(schema: dict[str, Any]) -> None:
        schema["type"] = "string"
        schema["description"] = description
        if pattern:
            schema["pattern"] = pattern
        schema.pop("title", None)

    return Field(pattern=pattern, json_schema_extra=_schema_extra)


def NumberField(
    description: str,
    *,
    minimum: float | None = None,
    maximum: float | None = None,
    exclusive_minimum: float | None = None,
    exclusive_maximum: float | None = None,
) -> Any:
    """Create a number field for form inputs.

    Args:
        description: Human-readable description shown as label.
        minimum: Optional minimum value (inclusive).
        maximum: Optional maximum value (inclusive).
        exclusive_minimum: Optional exclusive minimum value.
        exclusive_maximum: Optional exclusive maximum value.

    Example:
        ```python
        class MyForm(FormInput):
            age: float = NumberField(description="Your age", minimum=0, maximum=120)
            price: float = NumberField(
                description="Price in USD",
                minimum=0,
                exclusive_maximum=10000
            )
        ```
    """

    def _schema_extra(schema: dict[str, Any]) -> None:
        schema["type"] = "number"
        schema["description"] = description
        if minimum is not None:
            schema["minimum"] = minimum
        if maximum is not None:
            schema["maximum"] = maximum
        if exclusive_minimum is not None:
            schema["exclusiveMinimum"] = exclusive_minimum
        if exclusive_maximum is not None:
            schema["exclusiveMaximum"] = exclusive_maximum
        schema.pop("title", None)

    return Field(
        ge=minimum,
        le=maximum,
        gt=exclusive_minimum,
        lt=exclusive_maximum,
        json_schema_extra=_schema_extra,
    )


def DateTimeField(description: str) -> Any:
    """Create a datetime field for form inputs.

    Args:
        description: Human-readable description shown as label.

    Example:
        ```python
        from datetime import datetime

        class MyForm(FormInput):
            start_date: datetime = DateTimeField(description="Start date and time")
        ```
    """

    def _schema_extra(schema: dict[str, Any]) -> None:
        schema["type"] = "string"
        schema["format"] = "date-time"
        schema["description"] = description
        schema.pop("title", None)

    return Field(json_schema_extra=_schema_extra)


def DateField(description: str) -> Any:
    """Create a date (date only) field for form inputs.

    Args:
        description: Human-readable description shown as label.

    Example:
        ```python
        from datetime import date

        class MyForm(FormInput):
            start_date: date = DateField(description="Start date")
        ```
    """

    def _schema_extra(schema: dict[str, Any]) -> None:
        schema["type"] = "string"
        schema["format"] = "date"
        schema["description"] = description
        schema.pop("title", None)

    return Field(json_schema_extra=_schema_extra)


def FileField(
    description: str,
    *,
    multiple: bool = False,
) -> Any:
    """Create a file upload field for form inputs.

    The workflow receives signed URLs (strings) for the uploaded files.

    Args:
        description: Human-readable description shown as label.
        multiple: If True, allows uploading multiple files (use with ``list[str]`` type annotation).

    Example:
        ```python
        class MyForm(FormInput):
            document: str = FileField(description="Upload a document")
            attachments: list[str] = FileField(description="Upload files", multiple=True)
        ```
    """

    def _schema_extra(schema: dict[str, Any]) -> None:
        if multiple:
            schema["type"] = "array"
            schema["items"] = {"type": "string", "format": "uri", "title": "File"}
            schema["description"] = description
            schema.pop("title", None)
        else:
            schema["type"] = "string"
            schema["format"] = "uri"
            schema["title"] = "File"
            schema["description"] = description

    return Field(json_schema_extra=_schema_extra)


def SingleChoice(
    options: list[tuple[str, str]] | list[str],
    description: str | None = None,
) -> Any:
    """Create a single choice (dropdown/select) field for form inputs.

    Args:
        options: List of options, either:
            - List of strings: ["value1", "value2"] - value is used as both value and label
            - List of tuples: [("value1", "Label 1"), ("value2", "Label 2")]
        description: Optional overall description for the field.

    Returns:
        FieldInfo that produces oneOf JSON schema.

    Example:
        ```python
        class MyForm(FormInput):
            priority: str = SingleChoice(
                options=[
                    ("low", "Low Priority"),
                    ("medium", "Medium Priority"),
                    ("high", "High Priority"),
                ],
                description="Select the priority level"
            )
        ```
    """
    return Field(json_schema_extra=_SingleChoiceSchemaExtra(options, description))


def MultiChoice(
    options: list[tuple[str, str]] | list[str],
    description: str | None = None,
) -> Any:
    """Create a multi choice (multi-select) field for form inputs.

    Args:
        options: List of options, either:
            - List of strings: ["value1", "value2"] - value is used as both value and label
            - List of tuples: [("value1", "Label 1"), ("value2", "Label 2")]
        description: Optional overall description for the field.

    Returns:
        FieldInfo that produces array JSON schema with anyOf items.

    Example:
        ```python
        class MyForm(FormInput):
            tags: list[str] = MultiChoice(
                options=[
                    ("frontend", "Frontend"),
                    ("backend", "Backend"),
                    ("infra", "Infrastructure"),
                ],
                description="Select applicable tags"
            )
        ```
    """
    return Field(json_schema_extra=_MultiChoiceSchemaExtra(options, description))


class ConfirmationInputModel(BaseModel, Generic[ChoiceT]):
    model_config = ConfigDict(title="Confirmation", extra="forbid")
    choice: ChoiceT


def ConfirmationInput(
    options: list[tuple[ChoiceT, str]] | list[ChoiceT],
    description: str,
) -> type[ConfirmationInputModel[ChoiceT]]:
    """Create a ConfirmationInput schema for workflow confirmation dialogs.

    This creates a specialized input schema that renders as a confirmation UI
    in the workflow. The user must select one of the provided options.

    Args:
        options: List of (value, label) tuples for the confirmation options.
            Each tuple contains (value, label) where value is what you receive
            in the response and label is what the user sees.
        description: The description displayed to the user (mandatory).

    Returns:
        A ConfirmationInputModel class that can be used with wait_for_input().

    Example:
        ```python
        confirmation = await self.wait_for_input(
            ConfirmationInput(
                options=[
                    ("confirm", "Yes, proceed"),
                    ("cancel", "Cancel operation"),
                ],
                description="Do you want to continue with this action?"
            )
        )
        if confirmation.choice == "confirm":
            # proceed with action
            ...
        ```
    """
    if len(options) < 1:
        raise ValueError("Confirmation must have at least one option")

    schema_extra = _SingleChoiceSchemaExtra(options, description)

    class _ConfirmationInput(ConfirmationInputModel[ChoiceT]):
        choice: ChoiceT = Field(json_schema_extra=schema_extra)

        @model_validator(mode="after")
        def _validate_choice(self) -> Self:
            if self.choice not in schema_extra.allowed_values:
                raise ValueError(f"Invalid choice '{self.choice}'. Must be one of: {schema_extra.allowed_values}")
            return self

    return _ConfirmationInput


# Constants for first-class accept/decline support in the client
ACCEPT_OPTION_VALUE: Literal["accept"] = "accept"
DECLINE_OPTION_VALUE: Literal["decline"] = "decline"

AcceptChoice = Literal["accept"]
DeclineChoice = Literal["decline"]
AcceptDeclineChoice = Literal["accept", "decline"]


def is_accepted(confirmation: ConfirmationInputModel[ChoiceT]) -> bool:
    """Check if a confirmation result was accepted.

    Example:
        ```python
        confirmation = await self.wait_for_input(
            AcceptDeclineConfirmation(
                description="Do you want to proceed?",
            )
        )
        if is_accepted(confirmation):
            # proceed with action
            ...
        else:
            # user declined
            ...
        ```
    """
    return confirmation.choice == ACCEPT_OPTION_VALUE


@overload
def AcceptDeclineConfirmation(
    description: str,
    *,
    accept_label: str = "Accept",
    decline_label: None,
) -> type[ConfirmationInputModel[AcceptChoice]]: ...


@overload
def AcceptDeclineConfirmation(
    description: str,
    *,
    accept_label: str = "Accept",
    decline_label: str = "Decline",
) -> type[ConfirmationInputModel[AcceptDeclineChoice]]: ...


def AcceptDeclineConfirmation(
    description: str,
    *,
    accept_label: str = "Accept",
    decline_label: str | None = "Decline",
) -> type[ConfirmationInputModel[AcceptChoice]] | type[ConfirmationInputModel[AcceptDeclineChoice]]:
    """Create a confirmation dialog with accept/decline options that have keyboard hotkey support.

    This is a convenience wrapper around ConfirmationInput that uses special
    "accept" and "decline" values which the client recognizes for keyboard shortcuts.

    Args:
        description: The description displayed to the user (mandatory).
        accept_label: Label for the accept button. Defaults to "Accept".
        decline_label: Label for the decline button. Set to None for accept-only confirmation.
            Defaults to "Decline".

    Returns:
        A ConfirmationInputModel class that can be used with wait_for_input().
        The choice field will be "accept" or "decline".

    Example:
        ```python
        # Accept/Decline confirmation
        confirmation = await self.wait_for_input(
            AcceptDeclineConfirmation(
                description="Do you want to proceed with this action?",
                accept_label="Yes, proceed",
                decline_label="No, cancel",
            )
        )
        if is_accepted(confirmation):
            # proceed with action
            ...

        # Accept-only confirmation
        confirmation = await self.wait_for_input(
            AcceptDeclineConfirmation(
                description="Click to acknowledge and continue.",
                accept_label="I understand",
                decline_label=None,
            )
        )
        # User acknowledged, proceed...
        ```
    """
    if decline_label is None:
        options: list[tuple[AcceptChoice, str]] = [(ACCEPT_OPTION_VALUE, accept_label)]
        return ConfirmationInput(options=options, description=description)
    else:
        options_with_decline: list[tuple[AcceptDeclineChoice, str]] = [
            (DECLINE_OPTION_VALUE, decline_label),
            (ACCEPT_OPTION_VALUE, accept_label),
        ]
        return ConfirmationInput(options=options_with_decline, description=description)
