import weakref
from functools import partial
from typing import List

import omni.ui as ui
from omni.kit.property.usd.usd_attribute_model import UsdAttributeModel
from pxr import Sdf, Vt


class TokenArrayEditWidget:
    def __init__(self, stage, attr_name, prim_paths, metadata):
        self._model = UsdAttributeModel(stage, [path.AppendProperty(attr_name) for path in prim_paths], False, metadata)

        self._token_array_attrs = [stage.GetPrimAtPath(path).GetAttribute(attr_name) for path in prim_paths]

        self._frame = ui.Frame()
        self._frame.set_build_fn(self._build)

    def clean(self):
        self._model.clean()
        self._frame = None

    def _build(self):
        shared_tokens = None

        for attr in self._token_array_attrs:
            if attr.HasValue():
                tokens = list(attr.Get())
            else:
                tokens = []

            if shared_tokens is None:
                shared_tokens = tokens
            elif shared_tokens != tokens:
                shared_tokens = None
                break

        with ui.VStack(spacing=2):
            if shared_tokens is not None:
                for i, token in enumerate(shared_tokens):
                    with ui.HStack(spacing=2):
                        token_model = ui.StringField(name="models").model
                        token_model.set_value(token)

                        def on_edit_element(value_model, weak_self, index):
                            weak_self = weak_self()
                            if weak_self:
                                token_array = list(weak_self._token_array_attrs[0].Get())  # noqa: PLW0212
                                token_array[index] = value_model.as_string
                                weak_self._model.set_value(Vt.TokenArray(token_array))  # noqa: PLW0212

                        token_model.add_end_edit_fn(partial(on_edit_element, weak_self=weakref.ref(self), index=i))

                        def on_remove_element(weak_self, index):
                            weak_self = weak_self()
                            if weak_self:
                                token_array = list(weak_self._token_array_attrs[0].Get())  # noqa: PLW0212
                                token_array.pop(index)
                                weak_self._model.set_value(Vt.TokenArray(token_array))  # noqa: PLW0212

                        ui.Button(
                            "-",
                            width=ui.Pixel(14),
                            clicked_fn=partial(on_remove_element, weak_self=weakref.ref(self), index=i),
                        )

                def on_add_element(weak_self):
                    weak_self = weak_self()
                    if weak_self:
                        token_array_attr = weak_self._token_array_attrs[0]  # noqa: PLW0212
                        if token_array_attr.HasValue():
                            token_array = list(token_array_attr.Get())
                        else:
                            token_array = []

                        token_array.append("")
                        weak_self._model.set_value(Vt.TokenArray(token_array))  # noqa: PLW0212

                ui.Button(
                    "Add Element", width=ui.Pixel(30), clicked_fn=partial(on_add_element, weak_self=weakref.ref(self))
                )

            else:
                ui.StringField(name="models", read_only=True).model.set_value("Mixed")

    def _set_dirty(self):
        self._frame.rebuild()


def build_token_array_prop(
    stage,
    attr_name,
    metadata,
    property_type,
    prim_paths: List[Sdf.Path],
    additional_label_kwargs=None,
    additional_widget_kwargs=None,
):
    with ui.HStack(spacing=4):
        label_kwargs = {"name": "label", "width": 160, "height": 18}
        ui.Label(attr_name, **label_kwargs)
        ui.Spacer(width=5)
        return TokenArrayEditWidget(stage, attr_name, prim_paths, metadata)
