import pygame
from kivy.graphics.texture import Texture
from kivy.uix.widget import Widget

class PygameWidget(Widget):
    def __init__(self, width=800, height=600, **kwargs):
        super().__init__(**kwargs)
        self.width = width
        self.height = height

        pygame.init()
        self.surface = pygame.Surface((width, height))

        self.texture = Texture.create(size=(width, height))
        self.texture.flip_vertical()

        self._draw_callback = None
        self.bind(size=self._resize)

    def set_draw_callback(self, func):
        self._draw_callback = func

    def _resize(self, *args):
        w, h = self.size
        self.surface = pygame.Surface((int(w), int(h)))
        self.texture = Texture.create(size=(int(w), int(h)))
        self.texture.flip_vertical()

    def update(self, dt):
        if self._draw_callback:
            self._draw_callback(self.surface)

        data = pygame.image.tostring(self.surface, "RGBA", False)
        self.texture.blit_buffer(data, colorfmt="rgba", bufferfmt="ubyte")
        self.canvas.clear()
        with self.canvas:
            self.texture.blit(0, 0, self.width, self.height)
