import asyncio

from parat.services import ParaT, PipelineBuilder, PropertyManager
from paraview.modules.vtkRemotingPythonAsyncCore import (
    vtkPythonObservableWrapperUtilities,
)
import vtk
import sys

async def observeWebmChunks(view):

    frameCounter = -1
    maxNumFrames = 999

    async for chunk in vtkPythonObservableWrapperUtilities.GetIterator(view.GetWEBMChunkObservable()):

        if chunk:

            frameCounter += 1

            filename_suffix = f"{frameCounter}".zfill(len(str(maxNumFrames)))
            filename = f"frame-{filename_suffix}.bin"
            print(f"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("Done")


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

    builder = PipelineBuilder(session)

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

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

    await asyncio.sleep(0.1)

    view = await builder.CreateView("views", "RenderView")
    # Tells the view to use video codecs for higher compression ratio and lower latency.
    view.UseCompressedStreamsOn()
    # By turning on both, we can see the decompressed frame in the render window
    # and get a stream of webm chunks from the view.GetWEBMChunkObservable()
    view.DecodeStreamToRenderWindowAndWriteWebmChunks()
    chunkTask = asyncio.create_task(observeWebmChunks(view))

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

    await pmanager.Update(view)
    view.ResetCameraUsingVisiblePropBounds()

    if "-I" in sys.argv:
        iren = vtk.vtkRenderWindowInteractor()
        iren.SetRenderWindow(view.GetRenderWindow())
        iren.Initialize()
        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:
        await asyncio.sleep(0.1)

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

    # Must unsubscribe. otherwise, the stream is never closed and chunkTask waits forever.
    view.UnsubscribeWebmChunkStreams()

    await chunkTask

    # the view is destructed here, its webm writer and decoder get destroyed here.
    await App.finalize()


asyncio.run(main())
