import asyncio
from timeit import default_timer as timer
import vtk

import paraview.modules.vtkRemotingPythonAsyncCore as asyncCore
from parat.utils import *
import paraview.servermanager as sm

import os
import base64

try:
    import trame
except ModuleNotFoundError:
    print(
        "\n\nAre you running in a virtual enviroemnt with trame installed ? Check basicWebExample_README.md\n\n"
    )
    exit(0)


from trame import stop as stop_server
from trame import state, async_utils
from trame.layouts import SinglePage
from trame.html import vuetify


async def imageShow(viewProxy):
    """Use the image data from the view observable to update the web interface"""
    try:
        prev = None
        writer = vtk.vtkJPEGWriter()
        writer.SetQuality(80)
        writer.WriteToMemoryOn()
        i = 0
        with ProxyContextManager(viewProxy) as cm:
            async for value in cm.values():
                start = timer()
                writer.SetInputDataObject(value)
                writer.Write()
                image = writer.GetResult()
                imagebytes = memoryview(image).tobytes()
                encoded = base64.b64encode(imagebytes).decode("ascii")
                mime = "image/jpeg"
                state.image_url = f"data:{mime};base64,{encoded}"
                state.flush("image_url")
                end = timer()
                print(f"pushed {i} {end-start:.3f}")
                i += 1

            else:
                print("ImageTask Completed")

    except asyncio.CancelledError:
        print("ImageTask Cancelled")


RV = None
SPIN_Task = None
RUNNING = False


async def main():
    global RV, SPIN_Task
    await sm.initialize()

    filename = "disk_out_ref2.ex2"
    disk = sm.sources.IOSSReader()
    disk.FileName = filename
    disk.UpdateVTKObjects()

    rep = sm.rendering.GeometryRepresentation()
    rep.Input = [disk]
    rep.UpdateVTKObjects()

    # create a view
    rv = sm.CreateRenderView()
    rv.Representations = [rep]
    rv.UpdateVTKObjects()

    await rv.Update()
    rv.ResetCameraUsingVisiblePropBounds()
    await rv.Update()

    updateImageTask = async_utils.create_task(imageShow(rv))

    renderWindow = rv.GetRenderWindow()
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(renderWindow)

    iren.Initialize()
    RV = rv

    while RUNNING and not iren.GetDone():
        if SPIN_Task is None:
            iren.ProcessEvents()
            await asyncio.sleep(0.01)
        else:
            await rv.Update()

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

    rv.UnsubscribeImageStreams()
    RV = None

    await updateImageTask

    RV = None
    sm.finalize()


def main_helper():
    global RUNNING
    RUNNING = True
    task = async_utils.create_task(main())

    def info(result):
        print("Main task completed")
        RUNNING = False

    task.add_done_callback(info)


def spin_helper():
    global SPIN_Task
    if SPIN_Task is None:
        SPIN_Task = async_utils.create_task(spin())
    else:
        SPIN_Task.cancel()
        SPIN_Task = None


async def spin():
    rv = RV
    while True:
        rv.SMProxy.GetCamera().Azimuth(10)
        await asyncio.sleep(0.01)


# -----------------------------------------------------------------------------
# UI
# -----------------------------------------------------------------------------

layout = SinglePage("ImageStream")
layout.title.set_text("Async")

with layout.toolbar as tb:
    tb.dense = True
    vuetify.VSpacer()
    vuetify.VBtn("Start server", click=main_helper)
    vuetify.VBtn("Spin", click=spin_helper)
    vuetify.VBtn("Shutdown", click=stop_server)

with layout.content:
    vuetify.VImg(
        src=("image_url", None),
    )


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