"""Pull one RGB frame from a portal PipeWire ScreenCast stream."""
from PIL import Image


class PipeWireFrameBuffer:
    def __init__(self, fd, node_id):
        import gi
        gi.require_version('Gst', '1.0')
        from gi.repository import Gst
        Gst.init(None)
        self.Gst = Gst
        self.pipeline = Gst.parse_launch(
            f'pipewiresrc fd={int(fd)} path={int(node_id)} always-copy=true do-timestamp=true '
            'client-name=jarvis ! videoconvert ! video/x-raw,format=RGB ! '
            'appsink name=sink max-buffers=1 drop=true sync=false enable-last-sample=true'
        )
        self.sink = self.pipeline.get_by_name('sink')
        if not self.sink:
            raise RuntimeError('GStreamer appsink missing')
        self.pipeline.set_state(Gst.State.PLAYING)
        change, state, _pending = self.pipeline.get_state(5 * Gst.SECOND)
        if change == Gst.StateChangeReturn.FAILURE or state != Gst.State.PLAYING:
            self.close()
            raise RuntimeError('PipeWire ScreenCast pipeline failed to play')

    def capture(self, path, timeout_ms=4000):
        sample = self.sink.emit('try-pull-sample', timeout_ms * self.Gst.MSECOND)
        if sample is None:
            sample = self.sink.get_property('last-sample')
        if sample is None:
            raise RuntimeError('no PipeWire frame yet')
        buf = sample.get_buffer()
        caps = sample.get_caps().get_structure(0)
        width = int(caps.get_value('width'))
        height = int(caps.get_value('height'))
        ok, mapped = buf.map(self.Gst.MapFlags.READ)
        if not ok:
            raise RuntimeError('could not map PipeWire frame')
        try:
            data = bytes(mapped.data)
        finally:
            buf.unmap(mapped)
        stride = max(width * 3, len(data) // max(1, height))
        if stride != width * 3:
            rows = [data[i * stride:i * stride + width * 3] for i in range(height)]
            data = b''.join(rows)
        Image.frombytes('RGB', (width, height), data[:width * height * 3]).save(path)
        return path

    def close(self):
        if getattr(self, 'pipeline', None):
            try:
                self.pipeline.set_state(self.Gst.State.NULL)
            except Exception:
                pass
        self.pipeline = None
        self.sink = None