# python imports
import ctypes
# pip imports
import numpy as np
import wx
import skia


"""Enable high-res displays."""
try:
    ctypes.windll.shcore.SetProcessDpiAwareness(1)
    ctypes.windll.shcore.SetProcessDpiAwareness(2)
except Exception:
    pass  # fail on non-windows


class SkiaCPUCanvas(wx.Panel):
    """A cpu based skia canvas"""

    def __init__(self, parent, size):
        super().__init__(parent, size=size)
        # or else we'll have a flicker
        self.SetBackgroundStyle(wx.BG_STYLE_PAINT)

        self.Bind(wx.EVT_SIZE, self.on_size)
        self.Bind(wx.EVT_PAINT, self.on_paint)
        self.timer = wx.Timer(self)
        self.Bind(wx.EVT_TIMER, self.on_timer, self.timer)
        self.timer.Start(32)  # 1000/32 = 31.25 FPS

    def bgra_to_rgba(self, pixels: bytes, width: int, height: int) -> bytes:
        arr = np.frombuffer(pixels, dtype=np.uint8).reshape((height, width, 4))
        arr = arr[:, :, [2, 1, 0, 3]]  # Swap BGR -> RGB, keep Alpha
        return arr.tobytes()

    def on_paint(self, event):
        w, h = self.GetSize()
        if w == 0 or h == 0:
            return

        self.draw(w, h)

        image = self.surface.makeImageSnapshot()
        # image = image.toarray(colorType=skia.ColorType.kRGBA_8888_ColorType).tobytes()
        # pixels_rgba = self.bgra_to_rgba(image.tobytes(), w, h)
        bitmap = wx.Bitmap.FromBufferRGBA(w, h, image.tobytes())

        dc = wx.PaintDC(self)
        dc.DrawBitmap(bitmap, 0, 0)

    def draw(self, w, h):
        self.canvas.clear(skia.ColorBLACK)  # Clear the canvas with white background
        self.canvas.save()
        self.canvas.translate(w / 2, h / 2)

        paint = skia.Paint(AntiAlias=True, Color=skia.Color(0, 0, 255)) # BLUE COLOR
        self.canvas.drawRect(skia.Rect(0, 0, 100, 100), paint)
        self.canvas.restore()

    def on_size(self, event):
        wx.CallAfter(self.set_size)
        event.Skip()

    def set_size(self):
        # Actual drawing area (without borders) is GetClientSize
        size = self.GetClientSize()
        # On HiDPI screens, this could be 1.25, 1.5, 2.0, etc.
        scale = self.GetContentScaleFactor()
        width = int(size.width * scale)
        height = int(size.height * scale)
        self.size = wx.Size(width, height)

        self.surface = skia.Surface(width, height)
        self.canvas = self.surface.getCanvas()

        self.Refresh()

    def on_timer(self, event):
        self.Refresh()


class MainFrame(wx.Frame):
    def __init__(self):
        super().__init__(None, title="Skia CPU Canvas", size=(800, 600))
        panel = wx.Panel(self)
        sizer = wx.BoxSizer(wx.VERTICAL)
        self.canvas = SkiaCPUCanvas(panel, (800, 600))
        sizer.Add(self.canvas, 1, wx.EXPAND)
        panel.SetSizer(sizer)
        self.Show()


if __name__ == "__main__":
    app = wx.App(False)
    frame = MainFrame()
    frame.Show()
    app.MainLoop()

image = image.toarray(colorType=skia.ColorType.kRGBA_8888_ColorType).tobytes()
pixels_rgba = self.bgra_to_rgba(image.tobytes(), w, h)