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

from collections.abc import Callable, Iterator, Sequence
from typing import Any, Final

from marimo._output.formatters.structures import format_structure
from marimo._output.hypertext import Html
from marimo._output.rich_help import mddoc
from marimo._plugins.core.web_component import JSONType
from marimo._plugins.stateless.flex import hstack, vstack
from marimo._plugins.stateless.json_output import json_output
from marimo._plugins.ui._core.ui_element import UIElement


# Frontend type is a tuple (index, value update)
# Python type is a sequence of values, one for each UI element
@mddoc
class array(UIElement[dict[str, JSONType], Sequence[object]]):
    """An array of UI elements.

    Use an array to:
    - create a dynamic number of UI elements at runtime
    - group together logically related UI elements
    - keep the number of global variables in your program small

    Access the values of the elements using the `value` attribute of the array
    (`array.value`).

    The elements in the array can be accessed using square brackets
    (`array[index]`) and embedded in other marimo outputs. You can also
    iterate over the UI elements using the `in` operator (`for element in
    array`).

    Note: The UI elements in the array are clones of the original elements:
    interacting with the array will _not_ update the original elements, and
    vice versa.

    Examples:
        A heterogeneous collection of UI elements:

        ```python
        array = mo.ui.array([mo.ui.slider(1, 10), mo.ui.text(), mo.ui.date()])
        ```

        Get the values of the `slider`, `text`, and `date` elements via
        `array.value`:

        ```python
        # array.value returns a list with the values of the elements
        array.value
        ```

        Access and output a UI element in the array:

        ```python
        mo.md(f"This is a slider: array[0]")
        ```

        Some number of UI elements, determined at runtime:

        ```python
        mo.ui.array([mo.ui.slider(1, 10) for _ in range random.randint(4, 8)])
        ```

    Attributes:
        value (list): A list containing the values of the array's entries.
        elements (list): A list of the wrapped elements (clones of the originals).

    Args:
        elements (Sequence[UIElement]): The UI elements to include.
        label (str, optional): A descriptive name for the array. Defaults to "".
        on_change (Optional[Callable[[Sequence[object]], None]], optional): Optional callback
            to run when this element's value changes.
    """

    _name: Final[str] = "marimo-dict"

    def __init__(
        self,
        elements: Sequence[UIElement[Any, Any]],
        *,
        label: str = "",
        on_change: Callable[[Sequence[object]], None] | None = None,
    ) -> None:
        self._elements = [e._clone() for e in elements]
        self._label = label
        slotted_html = json_output(
            json_data=format_structure(self._elements),
            name="array" if not label else label,
        )
        super().__init__(
            component_name=array._name,
            initial_value={
                str(index): e._initial_value_frontend
                for index, e in enumerate(self._elements)
            },
            label=label,
            args={
                "element-ids": {
                    e._id: str(i) for i, e in enumerate(self._elements)
                },
            },
            slotted_html=slotted_html.text,
            on_change=on_change,
        )

        for i, element in enumerate(self._elements):
            element._register_as_view(self, key=str(i))

    @property
    def elements(self) -> Sequence[UIElement[JSONType, object]]:
        return self._elements

    def _convert_value(self, value: dict[str, JSONType]) -> Sequence[object]:
        if self._initialized:
            updating_value = getattr(self, "_updating_value", False)
            for k, v in value.items():
                element = self._elements[int(k)]
                # only call update if the value has changed
                if element._value_frontend != v:
                    if updating_value:
                        element._update_value(v)
                    else:
                        element._update(v)
        return [e._value for e in self._elements]

    def _on_update_completion(self) -> bool:
        any_updated = False
        for element in self._elements:
            any_updated |= element._on_update_completion()
        if any_updated:
            self._value = [e._value for e in self._elements]
        return any_updated

    def _clone(self) -> array:
        return array(
            elements=self.elements,
            label=self._label,
            on_change=self._on_change,
        )

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

    def __getitem__(self, key: int) -> UIElement[JSONType, object]:
        return self.elements[key]

    def __iter__(self) -> Iterator[UIElement[JSONType, object]]:
        return self.elements.__iter__()

    def __reversed__(self) -> Iterator[UIElement[JSONType, object]]:
        return self.elements.__reversed__()

    def __contains__(self, item: UIElement[JSONType, object]) -> bool:
        return item in self.elements

    def hstack(self, **kwargs: Any) -> Html:
        """Stack the elements horizontally.

        Args:
            **kwargs: Additional arguments passed to `marimo.hstack`.

        Returns:
            Html: HTML representation of horizontally stacked elements.
        """
        return hstack(items=self.elements, **kwargs)

    def vstack(self, **kwargs: Any) -> Html:
        """Stack the elements vertically.

        Args:
            **kwargs: Additional arguments passed to `marimo.vstack`.

        Returns:
            Html: HTML representation of vertically stacked elements.
        """
        return vstack(items=self.elements, **kwargs)
