import yaml
from trame_simput import get_simput_manager
from trame_simput.widgets import simput
from trame_simput.core.proxy import Proxy

from vtk import vtkObject

from parat.services import PropertyManager
import xml.etree.ElementTree as ET

QT_2_VUETIFY = {
    "root": "div",
    "group": "sw-group",
    "QCheckBox": "sw-switch",
    "pqDoubleRangeWidget": "sw-slider",
    "QComboBox": "sw-select",
    "pqDoubleLineEdit": "sw-text-field",  # FIXME
    "pqIntRangeWidget": "sw-slider",  # FIXME
    "ColorArray": None,  # FIXME
    "color_selector_with_palette": None,  # FIXME
}

# -----------------------------------------------------------------------------


class Resolver:
    def __init__(self):
        self._model = None
        self._labels = None

    def xml_elem_widget(self, entry, default=None):
        widget_name = entry.get("widget")
        if not widget_name:
            widget_name = default

        # print("widget_name", widget_name)

        if widget_name in QT_2_VUETIFY:
            ui_elem = QT_2_VUETIFY[widget_name]
            # print(" ===> ", ui_elem)
            if ui_elem:
                # print("~~~ OK ~~~")
                return ET.Element(ui_elem)
            return None

        # print("Unknown widget:", widget_name, entry)

    def xml_elem(self, name, property, label=None, help=None):
        ui_attributes = property.get("ui_attributes")
        xml_elem = self.xml_elem_widget(ui_attributes)
        if xml_elem is not None:
            # print("~~> add attr", name)
            xml_elem.set(":mtime", "data.mtime")
            xml_elem.set("name", name)
            xml_elem.set("label", label if label else name)
            xml_elem.set("size", f'{property.get("size", 1)}')
            xml_elem.set("type", property.get("type", "string"))
            xml_elem.set(":initial", f"data.original['{name}']")
            if help:
                xml_elem.set("help", help)
        else:
            print("~~> SKIP", xml_elem)

        return xml_elem

    def resolve(self, model, labels, ui=None):
        # print("resolve")
        groups = {}
        order = []
        last_group = -1
        for prop_name in model.get("_ordered_properties", []):
            property = model.get(prop_name)
            group_name = property.get("_group")

            if group_name is None:
                last_group = -1
                order.append(("property", prop_name))
            elif group_name != last_group:
                last_group = group_name
                order.append(("group", group_name))

            if group_name is not None:
                groups.setdefault(group_name, [])
                groups[group_name].append(prop_name)

        # Generate output
        out_root = ET.Element("div")
        for type, name in order:
            if type == "group":
                group_name = name
                # print("In group", name)
                group = model.get(f"_group_{group_name}")
                group_container = self.xml_elem_widget(group, "group")
                if group_container is None:
                    # print(" => skip group due to unkown widget", group)
                    continue

                out_root.append(group_container)
                group_container.set("title", group.get("label"))

                for prop_name in groups.get(group_name):
                    xml_elem = self.xml_elem(
                        prop_name,
                        model.get(prop_name),
                        labels.get(prop_name, {}).get("_label", prop_name),
                        labels.get(prop_name, {}).get("_help", None),
                    )
                    if xml_elem is not None:
                        # print("  - ", prop_name)
                        group_container.append(xml_elem)
                    # else:
                    #     print("  - (SKIPPED)", prop_name)

            if type == "property":
                prop_name = name
                xml_elem = self.xml_elem(
                    prop_name,
                    model.get(prop_name),
                    labels.get(prop_name, {}).get("_label", prop_name),
                    labels.get(prop_name, {}).get("_help", None),
                )
                if xml_elem is not None:
                    # print("Property (alone)", name)
                    out_root.append(xml_elem)
                # else:
                #     print("Property (alone-skipped)", name)

        # Serialize output XML as string
        # print("resolve -------------------")
        ET.indent(out_root, space="  ", level=0)
        # print(ET.tostring(out_root, encoding="utf-8").decode("utf-8"))
        return ET.tostring(out_root, encoding="utf-8").decode("utf-8")


# -----------------------------------------------------------------------------
class ObjectFactory:
    def __init__(self):
        self._next = None

    def next(self, proxy):
        self._next = proxy

    def create(self, name, **kwargs):
        obj = self._next
        self._next = None

        return obj


# -----------------------------------------------------------------------------
class ProxyAdapter:
    def __init__(self, server):
        self._server = server
        self.prop_manager = PropertyManager()

    def commit(self, simput_proxy):
        pv_proxy = simput_proxy.object
        for name in simput_proxy.edited_property_names:
            value = simput_proxy[name]
            if isinstance(value, Proxy):
                value = value.object if value else None
            elif value is None:
                continue

            self.prop_manager.SetValues(pv_proxy, **{name: value})

        self.prop_manager.Push(pv_proxy)
        self._server.controller.on_apply()

        return len(simput_proxy.edited_property_names)

    def reset(self, simput_proxy, *prop_name):
        pv_proxy = simput_proxy.object
        for name in prop_name:
            pv_property = pv_proxy.GetProperty(name)
            if pv_property is not None:
                pv_property.ClearUncheckedElements()

    def fetch(self, simput_proxy):
        pv_proxy = simput_proxy.object
        # await self.prop_manager.Fetch(pv_proxy)
        values = self.prop_manager.GetValues(
            pv_proxy, *simput_proxy.list_property_names()
        )
        for prop_name, prop_value in values.items():
            if isinstance(prop_value, vtkObject):
                print(
                    f"Need to properly handle property {prop_name} of type {prop_value.GetClassName()}"
                )
            else:
                simput_proxy.set_property(prop_name, prop_value)

        simput_proxy.commit()

    def update(self, simput_proxy, *property_names):
        pv_proxy = simput_proxy.object

        for name in property_names:
            value = simput_proxy[name]
            if isinstance(value, Proxy):
                value = value.object if value else None
            elif value is None:
                continue

            property = pv_proxy.GetProperty(name)

            if property.GetClassName() in [
                "vtkSMInputProperty",
                "vtkSMProxyProperty",
            ]:
                if isinstance(value, (list, tuple)):
                    for i, v in enumerate(value):
                        if isinstance(v, Proxy):
                            v = v.object if v else None
                        property.SetUncheckedProxy(i, v)
                else:
                    property.SetUncheckedProxy(0, value)
            elif isinstance(value, (list, tuple)):
                if len(value) != property.GetNumberOfUncheckedElements():
                    property.SetNumberOfUncheckedElements(len(value))
                for i, v in enumerate(value):
                    property.SetUncheckedElement(i, v)
            else:
                property.SetUncheckedElement(0, value)

    def before_delete(self, simput_proxy):
        pv_proxy = simput_proxy.object
        print(
            "Deleting PV proxy",
            simput_proxy.id,
            pv_proxy.GetGlobalID(),
            pv_proxy.GetReferenceCount(),
        )
        # # simple.Delete(pv_proxy)
        # print("simple.Delete() => done", pv_proxy.GetReferenceCount())


# -----------------------------------------------------------------------------


class ParaViewSimput:
    def __init__(self, server, name="pxm"):
        self._server = server
        self._factory = ObjectFactory()
        self._definitions = None
        self._simput = get_simput_manager(
            name,
            object_factory=self._factory,
            object_adapter=ProxyAdapter(self._server),
            resolver=Resolver(),
        )
        print(f"Create simput prefix={name}")
        self._root_widget = simput.Simput(
            self._simput, prefix=name, trame_server=server
        )
        self._pv_to_simput_ids = {}

    def to_sinput(self, pv_id_or_proxy):
        pv_id = None
        if isinstance(pv_id_or_proxy, (int, str)):
            # id
            pv_id = str(pv_id_or_proxy)
        else:
            # proxy
            pv_id = str(pv_id_or_proxy.GetGlobalID())

        simput_id = self._pv_to_simput_ids.get(pv_id, None)
        if simput_id is not None:
            return self._simput.proxymanager.get(simput_id)

        # Proxy not found
        return None

    def to_pv(self, simput_id_or_proxy):
        simput_proxy = simput_id_or_proxy
        if isinstance(simput_id_or_proxy, str):
            simput_proxy = self._simput.proxymanager.get(simput_id_or_proxy)

        return simput_proxy.object if simput_proxy else None

    def set_definition_manager(self, def_manager):
        self._definitions = def_manager

    @property
    def manager(self):
        return self._simput

    @property
    def pxm(self):
        return self._simput.proxymanager

    @property
    def root_widget(self):
        return self._root_widget

    def apply(self, **kwargs):
        self.root_widget.apply(**kwargs)

    def reset(self, **kwargs):
        self.root_widget.reset(**kwargs)

    async def create(self, pv_proxy):
        definition = self._definitions.GetDefinition(pv_proxy)
        self._simput.load_model(yaml_content=yaml.dump(definition.as_dict()))
        proxy_type = f"{pv_proxy.GetXMLGroup()}__{pv_proxy.GetXMLName()}"
        self._factory.next(pv_proxy)
        simput_proxy = self.pxm.create(proxy_type)
        self._pv_to_simput_ids[str(pv_proxy.GetGlobalID())] = simput_proxy.id
        simput_proxy.fetch()
        return simput_proxy
