from __future__ import annotations

import uuid
from typing import Any, Literal, Self

from mistralai_workflows.core.task.task import Task
from pydantic import BaseModel, ConfigDict, Field, model_validator


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

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


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 _AssistantMessageState(BaseModel):
    """Internal state for assistant message custom task events."""

    model_config = ConfigDict(title="assistant_message")
    contentChunks: list[TextOutput | ResourceOutput] = Field(default_factory=list)


async def send_assistant_message(text: str) -> 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: Text content to send.

    Example:
        ```python
        await send_assistant_message("Please provide your favorite Pokemon:")
        ```
    """
    async with Task[_AssistantMessageState](
        type="assistant_message",
        state=_AssistantMessageState(contentChunks=[TextOutput(text=text)]),
    ):
        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 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
    mimeType: Literal["application/vnd.mistral.canvas"]
    canvas: CanvasPayload


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


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


TextChunkType = Literal["text"]


class TextChunk(BaseModel):
    text: str
    type: TextChunkType = "text"


class ChatInputModel(BaseModel):
    model_config = ConfigDict(title="ChatInput", extra="forbid")
    message: list[TextChunk]


def ChatInput(prompt: str = "Type your message") -> type[ChatInputModel]:
    """Create a ChatInput schema with a custom prompt/placeholder.

    Args:
        prompt: The placeholder text shown in the chat input field.

    Example:
        ```python
        user_input = await self.wait_for_input(ChatInput("Ask a follow-up question"))
        ```
    """

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

    return _ChatInput


class _SingleChoiceSchemaExtra:
    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])
        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 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


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 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))
