import asyncio
import io
import math
import select
import sys
import termios
import tty

import pyvista as pv

import kitty


class TerminalPlotter(pv.Plotter):
    def __init__(self, width, height, **kwargs):
        super().__init__(off_screen=True, window_size=(width, height), **kwargs)
        self.width = width
        self.height = height
        self._running = True
        self._needs_render = True
        self._render_complete = False
        self._rendering = False
        self._render_pending = True 
        self.set_background([0.0, 0.0, 0.0])
        self.add_on_render_callback(self._on_render)
        # if we are too close to the bottom of the terminal, create some space.

    async def initialize(self):
        h_pix, _ = await kitty.get_terminal_cell_size()
        self.rows, _ = kitty.get_terminal_size()
        self.start_y, self.start_x = await kitty.get_position()
        self.needed_lines = math.ceil(self.height / h_pix)

        # Add vertical space if needed
        if self.rows - self.start_y < self.needed_lines:
            missing = self.needed_lines - (self.rows - self.start_y)
            self.start_y -= missing
            print("\n" * self.needed_lines, end="")
            kitty.set_position(self.start_y, self.start_x)


    async def render_to_kitty(self):
        self.render()
        buf = io.BytesIO()
        self.screenshot(buf, transparent_background=True, window_size=(self.width, self.height))
        await kitty.draw_to_terminal(buf)
        # print("y:", self.start_y, "rows:", self.rows, end="")
        kitty.set_position(self.start_y, self.start_x)

    async def _input_loop(self):
        fd = sys.stdin.fileno()
        old_settings = termios.tcgetattr(fd)
        try:
            tty.setcbreak(fd)
            while self._running:
                rlist, _, _ = select.select([sys.stdin], [], [], 0)
                if rlist:
                    c = sys.stdin.read(1)
                    if c == "q":  # quit on q
                        self._running = False
                    else:
                        await self._handle_key(c)
                await asyncio.sleep(0.01)
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)

    async def _handle_key(self, c):
        if c == "a":
            self.camera.Azimuth(10)
        elif c == "d":
            self.camera.Azimuth(-10)
        elif c == "w":
            self.camera.Elevation(10)
        elif c == "s":
            self.camera.Elevation(-10)
        self._needs_render = True
        # await self.render_to_kitty()

    def _on_render(self, *args):
        """Callback that runs after each render."""
        self._rendering = False
        self._render_complete = True

    async def _display_loop(self):
        """Loop that captures the screen and updates the terminal after rendering."""
        while self._running:
            if self._render_complete:
                self._render_complete = False
                buf = io.BytesIO()
                # Safe: Screenshot only after render is confirmed complete
                self.screenshot(buf, transparent_background=True, window_size=(self.width, self.height))
                await kitty.draw_to_terminal(buf)
                kitty.set_position(self.start_y, self.start_x)
            await asyncio.sleep(0.001)

    async def _render_loop(self):
        while self._running:
            if self._needs_render and not self._rendering:
                self._rendering = True
                self._needs_render = False
                self.update()  # triggers a render and calls _on_render
            await asyncio.sleep(0.001)

    async def _display_in_terminal(self, buf):
        await kitty.draw_to_terminal(buf)
        kitty.set_position(self.start_y, self.start_x)


    async def run(self):
        await self.initialize()
        self.iren.initialize()
        self._needs_render = True
        tasks = [
            asyncio.create_task(self._input_loop()),
            asyncio.create_task(self._render_loop()),
            asyncio.create_task(self._display_loop()),
        ]
        await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
        self._running = False


        #await self.render_to_kitty()
        #await self._input_loop()