# -----------------------------------------------------------------------------
# export PYTHONPATH=/home/seb/Documents/code/Async/build/lib/python3.10/site-packages
# export LD_LIBRARY_PATH=/home/seb/Documents/code/Async/build/lib
#
# python3.10 -m venv pv-venv
# source ./pv-venv/bin/activate
# pip install -U pip
# pip install trame
#
# python ./paraview/Remoting/Microservices/Testing/Trame/TestSCDemo.py
# or using pvpython:
# pvpython ./paraview/Remoting/Microservices/Testing/Trame/TestSCDemo.py --venv ./pv-venv
# -----------------------------------------------------------------------------
from wave import WAVE_FORMAT_PCM
from parat import venv
import asyncio
import base64
import json
from parat.services import (
    ApplyController,
    ParaT,
    PipelineBuilder,
    DefinitionManager,
    ActiveObjects,
    PropertyManager,
    PipelineViewer,
    ProgressObserver,
)
from parat.trame.simput import ParaViewSimput
from paraview.modules.vtkRemotingServerManager import vtkSMProxySelectionModel
from paraview.modules.vtkRemotingPythonAsyncCore import (
    vtkPythonObservableWrapperUtilities,
)

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

from trame.app import get_server, asynchronous
from trame.ui.vuetify import SinglePageWithDrawerLayout
from trame.widgets import vuetify, simput, trame, html, rca
from trame_rca.protocol import AreaAdapter

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

WAVELET_SCALAR_RANGE = [37.35310363769531, 276.8288269042969]


def generate_contour_values(data_range, number_of_contours):
    delta = (data_range[1] - data_range[0]) / (number_of_contours - 1)
    return [data_range[0] + (delta * float(i)) for i in range(number_of_contours)]


def apply_exec(**kwargs):
    print("Apply")


class App:
    def __init__(self, server=None):
        if server is None:
            server = get_server()

        self.server = server
        self.state = server.state
        self.ctrl = server.controller

        # ParaView Async
        self._app = ParaT()
        self._running = True
        self._ready = False
        self._pv_simput = ParaViewSimput(server, name="async")

        # internal state
        self.active_proxy = None
        self.active_representation = None
        self.active_view = None

        # state
        self.state.simput_active_source = 0
        self.state.spin_wait = 1 / 30  # target 30 fps
        self.state.wavelet_size = 10
        self.state.nb_contours = 10

        # controller
        self.ctrl.on_server_ready.add_task(self.initialize)
        self.ctrl.on_server_exited.add_task(self.finalize)
        self.ctrl.apply = self._pv_simput.apply
        self.ctrl.reset = self._pv_simput.reset
        self.ctrl.on_apply = apply_exec  # FIXME bug on trame-server

    async def initialize(self, **kwargs):
        self._session = await self._app.initialize()
        self._pipeline = PipelineViewer(self._session)
        self._builder = PipelineBuilder(self._session)
        self._def_mgt = DefinitionManager(self._session)
        self._active = ActiveObjects(self._session, "ActiveSources")
        self._prop_mgr = PropertyManager()
        self._progress = ProgressObserver(self._session)
        self._apply_ctrl = ApplyController(self._session)

        # Bind definition service to simput helper
        self._pv_simput.set_definition_manager(self._def_mgt)

        # default demo pipeline
        await self.setup_demo()

        # Tasks to monitor state change
        asynchronous.create_task(self.on_active_change())
        asynchronous.create_task(self.on_pipeline_change())
        asynchronous.create_task(self.monitor_server_status())
        asynchronous.create_task(self.monitor_view_stream())
        asynchronous.create_task(self.monitor_view_stats())
        asynchronous.create_task(self.monitor_progress_ds())
        asynchronous.create_task(self.monitor_progress_rs())

        # State listener
        self.state.change("view_size")(self.on_view_size_change)

        # RemoteControllerArea
        self._view_handler = AreaAdapter("view")
        self.ctrl.rc_area_register(self._view_handler)

        self._ready = True

    async def finalize(self, **kwargs):
        await self._app.finalize()

    async def monitor_server_status(self):
        while self._running:
            with self.state as state:
                await asyncio.sleep(state.spin_wait)

                # Spinning
                if state.spinning and self.active_view:
                    self.active_view.GetCamera().Azimuth(1)
                    self.active_view.StillRender()

                # Update client
                state.status_server += 5
                if state.status_server > 360:
                    state.status_server = 0

    async def on_active_change(self):
        async for proxy in self._active.GetCurrentObservable():
            self.active_proxy = proxy
            active_ids = []
            active_id = "0"

            if proxy:
                active_ids = [str(proxy.GetGlobalID())]
                simput_proxy = self._pv_simput.to_sinput(proxy)
                if simput_proxy is None:
                    simput_proxy = await self._pv_simput.create(proxy)
                active_id = simput_proxy.id
                await self._prop_mgr.UpdatePipeline(proxy)
                dataInformation = proxy.GetDataInformation(0)

            with self.state as state:
                state.git_tree_actives = active_ids
                self.state.simput_active_source = active_id

    async def on_pipeline_change(self):
        async for pipelineState in self._pipeline.GetObservable():
            list_to_fill = []
            for item in pipelineState:
                node = {
                    "name": item.GetName(),
                    "parent": str(
                        item.GetParentIDs()[0] if len(item.GetParentIDs()) > 0 else 0
                    ),
                    "id": str(item.GetID()),
                    "visibile": 1,
                }
                list_to_fill.append(node)

            with self.state as state:
                state.git_tree_sources = list_to_fill

    def ui_active_change(self, active):
        proxy = self._session.GetProxyManager().FindProxy(int(active[0]))
        self._active.SetCurrentProxy(proxy, vtkSMProxySelectionModel.CLEAR)

    async def create_proxy(self):
        with self.state as state:
            xml_group = state.xml_group
            xml_name = state.xml_name
            proxy = None
            if xml_group == "sources":
                proxy = await self._builder.CreateProxy(xml_group, xml_name)
            elif xml_group == "filters":
                input = self.active_proxy
                proxy = await self._builder.CreateProxy(
                    xml_group, xml_name, Input=input
                )
            elif xml_group == "representations":
                input = self.active_proxy
                view = self.active_view
                proxy = await self._builder.CreateRepresentation(
                    input, 0, view, xml_name
                )
                self.active_representation = proxy
            elif xml_group == "views":
                input = self.active_proxy
                proxy = await self._builder.CreateProxy(xml_group, xml_name)
                self.active_view = proxy
            else:
                print(f"Not sure what to create with {xml_group}::{xml_name}")

            # No proxy just skip work...
            if proxy is None:
                print("!!! No proxy created !!!")
                return

            # Load proxy definition
            simput_proxy = await self._pv_simput.create(proxy)

            if proxy == self.active_representation:
                state.simput_active_representation = simput_proxy.id

    async def setup_demo(self):
        view = await self._builder.CreateProxy("views", "RenderView")
        self.wavelet = await self._builder.CreateProxy("sources", "RTAnalyticSource")
        # self.clip = await self._builder.CreateProxy(
        #     "filters",
        #     "Clip",
        #     Input=self.wavelet,
        # )
        self.contour = await self._builder.CreateProxy(
            "filters",
            "Contour",
            Input=self.wavelet,
            ContourValues=generate_contour_values(WAVELET_SCALAR_RANGE, 10),
            SelectInputScalars=["", "", "", "", "RTData"],
        )

        # creating a representation before Apply() we avoid the default type
        # which is volume representation
        wave_representation = await self._builder.CreateRepresentation(
            self.wavelet, 0, view, "GeometryRepresentation"
        )

        self._prop_mgr.SetValues(wave_representation, Representation="Outline")

        self._prop_mgr.Push(self.contour, wave_representation)
        self._apply_ctrl.Apply()
        # gets the representation created in Apply()
        representation = await self._builder.CreateRepresentation(
            self.contour, 0, view, "GeometryRepresentation"
        )

        # View encoding setup
        self._prop_mgr.SetValues(
            view,
            force_push=True,
            LosslessMode=False,
            Display=False,
            StreamOutput=True,
        )

        # Keep track of view + rep
        self.active_view = view
        self.active_representation = representation

    def update_number_of_contours(self, nb_contours):
        if not self._ready:
            return

        self._prop_mgr.SetValues(
            self.contour,
            force_push=True,
            ContourValues=generate_contour_values(WAVELET_SCALAR_RANGE, nb_contours),
        )
        self._apply_ctrl.Apply()

    def update_wavelet_size(self, size):
        if not self._ready:
            return

        self._prop_mgr.SetValues(
            self.wavelet,
            force_push=True,
            WholeExtent=[-size, 0, -size, size, -size, size],
        )
        self._apply_ctrl.Apply()

    def on_view_size_change(self, view_size, **kwargs):
        size = view_size.get("size")
        self._prop_mgr.SetValues(
            self.active_view,
            ViewSize=(size.get("width"), size.get("height")),
            force_push=True,
        )
        self.active_view.StillRender()

    async def monitor_view_stream(self):
        async for package in vtkPythonObservableWrapperUtilities.GetIterator(
            self.active_view.GetViewOutputObservable()
        ):
            if package:
                self._view_handler.push(
                    package.GetData(), dict(type=package.GetMimeType())
                )
                # chunk = base64.b64encode(package.GetData()).decode("ascii")
                # mime = package.GetMimeType()
                # with self.state as state:
                #     state.img_url = f"data:{mime};base64,{chunk}"

    async def monitor_view_stats(self):
        async for metadata in vtkPythonObservableWrapperUtilities.GetIterator(
            self.active_view.GetViewStatsObservable()
        ):
            print(metadata)

    async def monitor_progress_ds(self):
        async for message in self._progress.GetServiceProgressObservable("ds"):
            data = json.loads(message)
            progress = data.get("ds", {}).get("Progress", 100)
            with self.state as state:
                state.status_ds_idle = progress == 100
                state.status_ds_progress = progress

    async def monitor_progress_rs(self):
        async for message in self._progress.GetServiceProgressObservable("rs"):
            data = json.loads(message)
            progress = data.get("rs", {}).get("Progress", 100)
            self.state.status_rs_idle = progress == 100


# -----------------------------------------------------------------------------
# Trame App
# -----------------------------------------------------------------------------

server = get_server()
app = App(server)


@app.state.change("wavelet_size")
def on_source_change(wavelet_size, **kwargs):
    print("wavelet_size", wavelet_size)
    app.update_wavelet_size(wavelet_size)


@app.state.change("nb_contours")
def on_contour_change(nb_contours, **kwargs):
    print("Update contours", nb_contours)
    app.update_number_of_contours(nb_contours)


# -----------------------------------------------------------------------------
# GUI
# -----------------------------------------------------------------------------

with SinglePageWithDrawerLayout(server) as layout:
    layout.root = app._pv_simput.root_widget

    layout.title.set_text("Parat")
    with layout.toolbar as toolbar:
        toolbar.dense = True
        vuetify.VSpacer()
        vuetify.VProgressCircular(
            "D",
            color="amber",
            size=35,
            width=5,
            indeterminate=("status_ds_idle", True),
            value=("status_ds_progress", 20),
            classes="mx-2",
        )
        vuetify.VProgressCircular(
            "R",
            color="purple",
            size=35,
            width=5,
            indeterminate=("status_rs_idle", True),
            value=("20",),
            classes="mx-2",
        )
        vuetify.VProgressCircular(
            "S",
            color="red",
            size=35,
            width=5,
            rotate=("status_server", 0),
            value=("20",),
            classes="mx-2",
        )
        vuetify.VProgressCircular(
            "C",
            color="teal",
            size=35,
            width=5,
            indeterminate=True,
            classes="mx-2",
        )
        vuetify.VDivider(vertical=True, classes="mx-2")
        vuetify.VCheckbox(
            small=True,
            v_model=("spinning", False),
            dense=True,
            classes="mx-2",
            hide_details=True,
            on_icon="mdi-axis-z-rotate-counterclockwise",
            off_icon="mdi-axis-z-rotate-counterclockwise",
        )
        with vuetify.VBtn(
            icon=True, small=True, click=app.ctrl.view_reset_camera, classes="mx-2"
        ):
            vuetify.VIcon("mdi-crop-free")

    with layout.drawer as drawer:
        drawer.width = 300
        trame.GitTree(
            sources=("git_tree_sources", []),
            actives=("git_tree_actives", []),
            actives_change=(app.ui_active_change, "[$event]"),
        )
        with vuetify.VCard(classes="mb-2 mx-1"):
            vuetify.VCardTitle("Wavelet {{ wavelet_size }}", classes="py-0")
            vuetify.VDivider()
            with vuetify.VCardText():
                vuetify.VSlider(
                    v_model=("wavelet_size", 10),
                    min=10,
                    max=1000,
                    step=10,
                    hide_details=True,
                    dense=True,
                )
        with vuetify.VCard(classes="mb-2 mx-1"):
            vuetify.VCardTitle("Contours {{ nb_contours }}", classes="py-0")
            vuetify.VDivider()
            with vuetify.VCardText():
                vuetify.VSlider(
                    v_model=("nb_contours", 10),
                    min=10,
                    max=1000,
                    step=10,
                    hide_details=True,
                    dense=True,
                )
        # simput.SimputItem(item_id=("simput_active_source", None))

    with layout.content:
        with vuetify.VContainer(fluid=True, classes="pa-0 fill-height"):
            rca.RemoteControlledArea(
                name="view",
            )
            # with trame.SizeObserver("view_size"):
            #     html.Img(
            #         src=("img_url", ""),
            #         style="position: absolute; left: 0; top: 0;",
            #     )

# -----------------------------------------------------------------------------
# CLI
# -----------------------------------------------------------------------------

if __name__ == "__main__":
    server.start()
