FUNCTIONS/json_template_drafting_functions.py

import marimo as mo
import json


def create_widgets(widget_type, values, widget_dict=None, stack=None):
    """
    Create marimo widgets for a list of values.

    Args:
        widget_type: Default marimo widget type (e.g., mo.ui.text_area)
        values: List of values to create widgets for
        widget_dict: Optional dict mapping specific values to widget types
        stack: 'horizontal' for hstack, 'vertical' for vstack, None for dict return
    """
    widgets = {}

    for value in values:
        # Use specific widget type if provided, otherwise use default
        current_widget_type = (
            widget_dict.get(value, widget_type) if widget_dict else widget_type
        )
        widgets[value] = current_widget_type(label=f"**{str(value)}**")

    if stack == "horizontal":
        return mo.hstack(list(widgets.values()), gap=1)
    elif stack == "vertical":
        return mo.vstack(list(widgets.values()), gap=1)
    else:
        return widgets


def json_to_marimo_ui(
    data,
    replicate_values=False,
    label="",
    label_containers=False,
    label_nested_leaves=False,
):
    """
    Convert a JSON structure to marimo UI elements.

    Args:
        data: JSON-like dict/list, or mo.ui.file element containing JSON
        replicate_values: If True, populate widgets with existing values; if False, leave empty
        label: Optional label for the root element
        label_containers: If True, apply labels to nested dicts/arrays (causes double nesting display);
                            if False (default), only label leaf input widgets
        label_nested_leaves: If True, label leaf widgets with their key names;
                            if False (default), only root label is applied, nested leaves have no labels

    Returns:
        mo.ui.dictionary for dicts, mo.ui.array for lists
    """
    # import marimo as mo
    # import json

    # Handle mo.ui.file input
    if hasattr(data, "contents") and callable(data.contents):
        file_contents = data.contents()
        if file_contents is None:
            return mo.ui.dictionary({}, label=label)
        if isinstance(file_contents, bytes):
            file_contents = file_contents.decode("utf-8")
        data = json.loads(file_contents)

    def _convert(obj, replicate, lbl="", is_root=False):
        if isinstance(obj, dict):
            elements = {
                str(key): _convert(
                    value,
                    replicate,
                    lbl=str(key) if label_nested_leaves else "",
                )
                for key, value in obj.items()
            }
            container_label = lbl if (is_root or label_containers) else ""
            return mo.ui.dictionary(elements, label=container_label)

        elif isinstance(obj, list):
            if not obj:
                container_label = lbl if (is_root or label_containers) else ""
                return mo.ui.array([], label=container_label)
            elements = [
                _convert(item, replicate, lbl="" if not label_nested_leaves else "")
                for item in obj
            ]
            container_label = lbl if (is_root or label_containers) else ""
            return mo.ui.array(elements, label=container_label)

        elif isinstance(obj, bool):
            leaf_label = lbl if (is_root or label_nested_leaves) else ""
            return mo.ui.checkbox(
                value=obj if replicate else False,
                label=leaf_label,
            )

        elif isinstance(obj, int):
            leaf_label = lbl if (is_root or label_nested_leaves) else ""
            return mo.ui.number(
                value=obj if replicate else None,
                label=leaf_label,
                full_width=True,
            )

        elif isinstance(obj, float):
            leaf_label = lbl if (is_root or label_nested_leaves) else ""
            return mo.ui.number(
                value=obj if replicate else None,
                step=0.01,
                label=leaf_label,
                full_width=True,
            )

        elif isinstance(obj, str):
            val = obj if replicate else ""
            leaf_label = lbl if (is_root or label_nested_leaves) else ""
            if len(obj) > 40:
                return mo.ui.text_area(value=val, label=leaf_label, full_width=True)
            return mo.ui.text(value=val, label=leaf_label, full_width=True)

        elif obj is None:
            leaf_label = lbl if (is_root or label_nested_leaves) else ""
            return mo.ui.text(value="", label=leaf_label, full_width=True)

        else:
            leaf_label = lbl if (is_root or label_nested_leaves) else ""
            return mo.ui.text(
                value=str(obj) if replicate else "",
                label=leaf_label,
                full_width=True,
            )

    return _convert(data, replicate_values, label, is_root=True)


def create_download_button(
    widget_value, filename_prefix="download", add_uuid_suffix=True
):
    """Create a marimo download button appropriate for the widget value type.

    Automatically detects data type and creates download button with correct format and MIME type.

    Args:
        widget_value: Data to download (dict/list as JSON, DataFrames as CSV,
            str as text, bytes as binary, others converted to string)
        filename_prefix (str): Filename prefix. Defaults to "download".
        add_uuid_suffix (bool): Add 4-char UUID suffix to filename. Defaults to True.

    Returns:
        mo.download: Configured marimo download button widget.

    Examples:
        >>> data = {"key": "value", "numbers": [1, 2, 3]}
        >>> button = create_download_button(data, "my_data")

        >>> df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
        >>> button = create_download_button(df, "dataframe")

    Note:
        Requires marimo (mo) and uuid modules. JSON formatted with 2-space indent.
        CSV exports exclude index for pandas DataFrames.
    """
    import marimo as mo
    import json
    import uuid

    if isinstance(widget_value, dict) or isinstance(widget_value, list):
        # JSON data
        data = json.dumps(widget_value, indent=2)
        filename = f"{filename_prefix}.json"
        if add_uuid_suffix:
            uuid_suffix = str(uuid.uuid4())[:4]
            filename = f"{filename_prefix}_{uuid_suffix}.json"
        return mo.download(
            data=data.encode(), filename=filename, mimetype="application/json"
        )

    elif hasattr(widget_value, "to_csv"):
        # Pandas DataFrame
        data = widget_value.to_csv(index=False)
        filename = f"{filename_prefix}.csv"
        if add_uuid_suffix:
            uuid_suffix = str(uuid.uuid4())[:4]
            filename = f"{filename_prefix}_{uuid_suffix}.csv"
        return mo.download(data=data.encode(), filename=filename, mimetype="text/csv")

    elif hasattr(widget_value, "write_csv"):
        # Polars DataFrame
        data = widget_value.write_csv()
        filename = f"{filename_prefix}.csv"
        if add_uuid_suffix:
            uuid_suffix = str(uuid.uuid4())[:4]
            filename = f"{filename_prefix}_{uuid_suffix}.csv"
        return mo.download(data=data.encode(), filename=filename, mimetype="text/csv")

    elif isinstance(widget_value, str):
        # Text data
        filename = f"{filename_prefix}.txt"
        if add_uuid_suffix:
            uuid_suffix = str(uuid.uuid4())[:4]
            filename = f"{filename_prefix}_{uuid_suffix}.txt"
        return mo.download(
            data=widget_value.encode(),
            filename=filename,
            mimetype="text/plain",
        )

    elif isinstance(widget_value, bytes):
        # Binary data
        filename = f"{filename_prefix}.bin"
        if add_uuid_suffix:
            uuid_suffix = str(uuid.uuid4())[:4]
            filename = f"{filename_prefix}_{uuid_suffix}.bin"
        return mo.download(
            data=widget_value,
            filename=filename,
            mimetype="application/octet-stream",
        )

    else:
        # Fallback: convert to string
        data = str(widget_value)
        filename = f"{filename_prefix}.txt"
        if add_uuid_suffix:
            uuid_suffix = str(uuid.uuid4())[:4]
            filename = f"{filename_prefix}_{uuid_suffix}.txt"
        return mo.download(data=data.encode(), filename=filename, mimetype="text/plain")


if __name__ == "__main__":
    print("Only Importable in Marimo Notebooks")