import asyncio

from parat.services import ParaT, PipelineBuilder, PropertyManager, ProgressObserver, ReactiveCommandHelper
from paraview.modules.vtkRemotingPythonAsyncCore import vtkPythonObservableWrapperUtilities
import vtk
import sys

cond = asyncio.Condition()

async def monitorViewOutput(view):

    global cond
    frameCounter = -1
    maxNumFrames = 999999

    async for package in vtkPythonObservableWrapperUtilities.GetIterator(
        view.GetViewOutputObservable()
    ):

        if package:
            chunk = package.GetData()
            frameCounter += 1
            async with cond:
                cond.notify()
            filename_suffix = f"{frameCounter}".zfill(len(str(maxNumFrames)))
            filename = f"frame-{filename_suffix}.bin"
            print(f"==> {filename_suffix}: write-chunk {filename}")

            with open(filename, "wb") as f:
                # prepare the bytes string from chunk's values
                as_bytes = b"".join(
                    [
                        chunk.GetValue(i).to_bytes(1, sys.byteorder)
                        for i in range(chunk.GetNumberOfValues())
                    ]
                )
                f.write(as_bytes)
        else:
            continue

    print("monitorViewOutput Done")


async def monitorViewStats(progress):

    frameCounter = -1
    maxNumFrames = 999999

    async for stats in progress.GetServiceMetricsObservable("rs"):
        frameCounter += 1
        frameIdPadded = f"{frameCounter}".zfill(len(str(maxNumFrames)))
        print(f"==> {frameIdPadded}: {stats}")

    print("monitorViewStats Done")

async def main():
    App = ParaT()
    session = await App.initialize()

    builder = PipelineBuilder(session)

    sphere = await builder.CreateProxy(
        "sources", "SphereSource", ThetaResolution=80, Radius=2
    )
    shrink = await builder.CreateProxy(
        "filters", "ShrinkFilter", Input=sphere, ShrinkFactor=0.3
    )

    pmanager = PropertyManager()
    pmanager.SetValues(sphere, Radius=3, Center=[1, 2, 3], EndPhi=90)

    progress = ProgressObserver(session)

    await asyncio.sleep(0.1)

    view = await builder.CreateProxy("views", "RenderView")
    chunkTask = asyncio.create_task(monitorViewOutput(view))
    statsTask = asyncio.create_task(monitorViewStats(progress))
    # Set range of stream quality
    pmanager.SetValues(view, InteractiveStreamQuality=60)
    pmanager.SetValues(view, StillStreamQuality=90)
    # Disable lossless
    pmanager.SetValues(view, LosslessMode=False)
    # Disable jpeg
    pmanager.SetValues(view, JpegMode=False)
    # Use Software encoder, defaults to VP9
    pmanager.SetValues(view, VideoEncoderType='VpxEncoder')
    # Multiplex the stream by setting this to true. we're interested in webm and not raw vp9
    pmanager.SetValues(view, MediaSegmentsMode=False)
    # Enable client side decode. You will want to disable display when MediaSegmentsMode=True
    pmanager.SetValues(view, Display=True)
    # Enable output stream.
    # When codec type is lossless, interaction will slow down due to I/O bottleneck in chunkTask.
    pmanager.SetValues(view, StreamOutput=True)

    representation = await builder.CreateRepresentation(
        shrink, 0, view, "GeometryRepresentation"
    )

    # before we 'update' the view, let's install an interactor on the view
    # so that the viewport dimensions, cameras of the client and server's render windows sync up.
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(view.GetRenderWindow())

    pmanager.Push(view)
    await pmanager.Update(view)
    if "-I" in sys.argv:
        while not iren.GetDone():
            iren.ProcessEvents()
            await asyncio.sleep(0.01)
            if iren.GetKeyCode() == "s":
                pmanager.SetVisibility(representation, view, True)
            if iren.GetKeyCode() == "h":
                pmanager.SetVisibility(representation, view, False)
    else:
        cam = iren.GetRenderWindow().GetRenderers().GetFirstRenderer().GetActiveCamera()
        for i in range(5):
            view.GetCamera().Azimuth(10)
            view.StillRender()
            async with cond:
                await cond.wait()

    status = await builder.DeleteProxy(shrink)
    assert status == True

    view.UnsubscribeOutputStreams()
    await chunkTask

    progress.Unsubscribe()
    await statsTask

    # the view is destructed here.
    # on its way out, the output stream subscribers are unsubbed here,
    # the webm writer + decoder get finalized here.
    await App.finalize()


asyncio.run(main())
