from PySide6.QtCore import QObject, QUrl, Qt, Signal, Slot
from PySide6.QtGui import QGuiApplication, QOffscreenSurface, QOpenGLContext, QSurfaceFormat
from PySide6.QtOpenGL import QOpenGLFramebufferObject
from PySide6.QtQml import QQmlComponent, QQmlEngine, QmlElement
from PySide6.QtQuick import QQuickGraphicsDevice, QQuickRenderControl, QQuickRenderTarget, QQuickWindow, QSGRendererInterface
from PySide6.QtWebEngineQuick import QtWebEngineQuick

import sys

QML_IMPORT_NAME = "OffscreenTest.PageCapture"
QML_IMPORT_MAJOR_VERSION = 1
QML_IMPORT_MINOR_VERSION = 0

@QmlElement
class PageCapture(QObject):
    def __init__(self, parent = None):
        super().__init__(parent)
        self.control = None
        self.fb = None

    @Slot()
    def screenshot(self):
        if self.control is None or self.fb is None:
            return

        self.control.polishItems()
        self.control.beginFrame()
        self.control.sync()
        self.control.render()
        self.control.endFrame()

        self.fb.toImage().save("offscreen.png")

def on_status_changed(component, window, control, fb):
    root = component.create()
    if component.isError():
        for e in component.errors():
            print(e.url(), e.line(), e)
        return
    root.setParentItem(window.contentItem())
    capture = root.findChild(PageCapture, "capture")
    capture.control = control
    capture.fb = fb

def main(argv=None):
    QtWebEngineQuick.initialize()
    QGuiApplication.setAttribute(Qt.AA_ShareOpenGLContexts)
    app = QGuiApplication(argv)

    QQuickWindow.setGraphicsApi(QSGRendererInterface.OpenGL)

    engine = QQmlEngine()
    component = QQmlComponent(engine)

    fmt = QSurfaceFormat()
    fmt.setMajorVersion(4)
    fmt.setMinorVersion(6)
    fmt.setDepthBufferSize(16)
    fmt.setStencilBufferSize(8)

    context = QOpenGLContext()
    context.setFormat(fmt)
    context.create()
    if not context.isValid():
        return

    surface = QOffscreenSurface()
    surface.setFormat(fmt)
    surface.create()
    if not surface.isValid():
        return

    context.makeCurrent(surface)

    control = QQuickRenderControl()
    window = QQuickWindow(control)
    window.setGraphicsDevice(QQuickGraphicsDevice.fromOpenGLContext(context))
    window.setGeometry(0, 0, 640, 480)

    if not control.initialize():
        return

    fb = QOpenGLFramebufferObject(640, 480, QOpenGLFramebufferObject.CombinedDepthStencil)

    target = QQuickRenderTarget.fromOpenGLTexture(fb.texture(), fb.size())
    window.setRenderTarget(target)

    component.statusChanged.connect(lambda: on_status_changed(component, window, control, fb))

    component.loadUrl(QUrl.fromLocalFile("offscreen.qml"))

    engine.quit.connect(app.quit)

    app.exec()

if __name__ == "__main__":
    main(sys.argv)