from typing import Callable
import weakref
import carb
import omni.timeline
import omni.ui as ui
import omni.usd

from pxr import Usd, Sdf, Tf, Gf

from omni.kit.property.usd.usd_attribute_model import (
    FloatModel,
    UsdBase,
    UsdAttributeModel,
    GfVecAttributeModel,
)

# TODO: will add controlstate

class VecAttributeModel(ui.AbstractItemModel):
    def __init__(self, default_value, begin_edit_callback: Callable[[None], None] = None, end_edit_callback: Callable[[Gf.Vec3d], None] = None):
        super().__init__()

        # We want to parse the values of this model to look for math operations,
        # so its items need to use string models instead of float models
        class StringModel(ui.SimpleStringModel):
            def __init__(self, parent, index):
                super().__init__()
                self._parent = weakref.ref(parent)
                self.index = index
            def begin_edit(self):
                parent = self._parent()
                parent.begin_edit(self)
            def end_edit(self):
                parent = self._parent()
                parent.end_edit(self)

        class VectorItem(ui.AbstractItem):
            def __init__(self, model):
                super().__init__()
                self.model = model

        dimension = 3
        self._items = [VectorItem(StringModel(self, i)) for i in range(dimension)]
        self._default_value = default_value
        self._begin_edit_callback = begin_edit_callback
        self._end_edit_callback = end_edit_callback

        for item in self._items:
            item.model.set_value(self._default_value)

        self._root_model = ui.SimpleStringModel()
        self._root_model.add_value_changed_fn(lambda a: self._item_changed(None))

    def get_item_children(self, item):
        return self._items

    def get_item_value_model(self, item, column_id):
        if item is None:
            return self._root_model
        return item.model

    def _on_value_changed(self, item):
        pass

    def begin_edit(self, item):
        index = item.index
        if self._begin_edit_callback:
            self._begin_edit_callback(index)

    def end_edit(self, model):
        text = model.get_value_as_string()
        index = model.index

        for item in self._items:
            item.model.set_value(self._default_value)

        if self._end_edit_callback:
            self._end_edit_callback(text, index)
