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

import asyncio
import io
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, cast

import marimo._output.data.data as mo_data
from marimo._output.rich_help import mddoc
from marimo._plugins.core.media import (
    guess_mime_type,
    io_to_data_url,
    is_data_empty,
    mime_type_to_ext,
)
from marimo._plugins.ui._core.ui_element import UIElement
from marimo._runtime.functions import EmptyArgs, Function

if TYPE_CHECKING:
    from collections.abc import Callable, Coroutine

DataType = str | bytes | io.BytesIO | io.BufferedReader


@dataclass
class LoadResponse:
    data: str
    filename: str | None


@mddoc
class download(UIElement[None, None]):
    """
    Show a download button for a url, bytes, or file-like object.

    Args:
        data (Union[str, bytes, io.BytesIO, callable]): The data to download. Can be:
            - string (interpreted as a URL)
            - bytes
            - file opened in binary mode
            - callable returning any of the above (for lazy loading)
            - async callable returning any of the above (for lazy loading)
        filename (str, callable): The name of the file to download. Can be a
            zero-arg callable returning a string, evaluated at click time so
            the name reflects the latest application state. A callable filename
            disables extension-based mimetype inference; pass `mimetype`
            explicitly if you need a specific one.
            If not provided, the name will be guessed from the data.
        mimetype (str): The mimetype of the file to download, for example,
            (e.g. "text/csv", "image/png"). If not provided,
            the mimetype will be guessed from the filename.
        disabled (bool): Whether to disable the download button.
        label (str): The label of the download button.

    Example:
        ```python
        # Eager loading
        download_txt = mo.download(
            data="Hello, world!".encode("utf-8"),
            filename="hello.txt",
            mimetype="text/plain",
        )


        # Lazy loading
        def get_large_data():
            return b"large data"


        download_lazy = mo.download(
            data=get_large_data,
            filename="large.txt",
        )
        ```
    """

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

    def __init__(
        self,
        data: DataType
        | Callable[[], DataType]
        | Callable[[], Coroutine[None, None, DataType]],
        filename: str | Callable[[], str] | None = None,
        mimetype: str | None = None,
        disabled: bool = False,
        *,
        label: str = "Download",
    ) -> None:
        self._data = data
        self._filename = filename
        self._mimetype = mimetype

        data_url = ""
        is_lazy = callable(data) or callable(filename)

        # name used to guess mimetype; a callable filename has no extension to
        # inspect at render time, so skip inference and fall back to text/plain
        name_for_mime = (
            data
            if isinstance(data, str)
            else (None if callable(filename) else filename)
        )
        resolved_mimetype = (
            mimetype or guess_mime_type(name_for_mime) or "text/plain"
        )
        ext = mime_type_to_ext(resolved_mimetype) or "txt"

        # Convert to bytes right away since can only be read once
        if isinstance(data, io.BufferedReader):
            filename = filename or data.name
            data.seek(0)
            data = data.read()

        # When non-lazy
        if not callable(data):
            # Maybe update disabled
            disabled = disabled or is_data_empty(data)

            # Maybe update name
            if filename is None and hasattr(data, "name"):
                filename = cast(str, cast(Any, data).name)

            # create a virtual file to avoid loading the data in the browser
            # only if the data is not lazy
            data_url = mo_data.any_data(data, ext=ext).url

        super().__init__(
            component_name=self._name,
            initial_value=None,
            label=label,
            on_change=None,
            args={
                "data": data_url,
                "filename": None if callable(filename) else filename,
                "disabled": disabled,
                "lazy": is_lazy,
            },
            functions=(
                (
                    Function(
                        name="load",
                        arg_cls=EmptyArgs,
                        function=self._load,
                    ),
                )
            ),
        )

    async def _load(self, _args: EmptyArgs) -> LoadResponse:
        filename = (
            self._filename() if callable(self._filename) else self._filename
        )

        # Eager data already ships as the button's href; a callable filename is
        # the only reason load runs here, so resolve the name and let the
        # frontend reuse the href instead of re-encoding the payload.
        if not callable(self._data):
            return LoadResponse(data="", filename=filename)

        if isinstance(self._data, UIElement):
            result = self._data
        else:
            result_or_coroutine = self._data()
            if asyncio.iscoroutine(result_or_coroutine):
                result = await result_or_coroutine
            else:
                result = result_or_coroutine

        url = io_to_data_url(
            result, fallback_mime_type=self._mimetype or "text/plain"
        )
        if url is None:
            raise ValueError("Failed to convert data to data URL")
        return LoadResponse(data=url, filename=filename)

    def _convert_value(self, value: None) -> None:
        return value
