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 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):
    outputs: list[TextOutput | ResourceOutput] = Field(default_factory=list)


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