import pyvista as pv
import vtk
from vtk.util import numpy_support
import numpy as np
import time
import kitty
import asyncio
from PIL import Image
from io import BytesIO


async def render_to(image_array):
    """
    Dummy rendering function that receives the rendered image as a numpy array.
    Replace this with your actual rendering/display logic.
    """
    # print(f"Received image with shape: {image_array.shape} and dtype: {image_array.dtype}")
    y, x = await kitty.get_position()
    img = Image.fromarray(image_array)
    buffer: BytesIO = BytesIO()
    img.save(buffer, format="PNG")
    await kitty.draw_to_terminal(buffer)
    kitty.set_position(y, x)


async def main():
    plotter = pv.Plotter(off_screen=True, window_size=[400, 400])
    plotter.add_mesh(pv.Sphere(), color="tomato")
    plotter.add_text("Live capture example", font_size=14)

    # plotter.show(auto_close=False)  # keep window open and interactive

    # Access the vtkRenderWindow
    vtk_render_window = plotter.ren_win

    # Setup the window-to-image filter to capture the rendered scene
    w2i_filter = vtk.vtkWindowToImageFilter()
    w2i_filter.SetInput(vtk_render_window)

    try:
        while True:
            # Trigger a render to make sure the content is fresh
            vtk_render_window.Render()

            # Update the filter to grab the current buffer
            w2i_filter.Modified()
            w2i_filter.Update()

            vtk_image = w2i_filter.GetOutput()

            width, height, _ = vtk_image.GetDimensions()
            vtk_array = vtk_image.GetPointData().GetScalars()
            arr = numpy_support.vtk_to_numpy(vtk_array)

            # Reshape the array to height x width x channels (probably 3 or 4)
            arr = arr.reshape(height, width, -1)

            # Flip vertically because VTK's origin is bottom-left
            arr = np.flip(arr, axis=0)

            # Call the dummy render_to function with captured frame
            await render_to(arr)

            time.sleep(0.1)  # adjust framerate to your needs, e.g. 10 FPS
            plotter.camera.Azimuth(1)

    except KeyboardInterrupt:
        print("Capture stopped by user.")
        plotter.close()


if __name__ == "__main__":
    asyncio.run(main())