import importlib

import numpy as np

import pyqtgraph as pg
import pyqtgraph.opengl as gl
from pyqtgraph.Qt import QtCore, QtGui
from pyqtgraph.examples.utils import FrameCounter

pngfile = importlib.resources.files("pyqtgraph.icons.peegee") / "peegee_512px@2x.png"
qimg = QtGui.QImage.fromData(pngfile.read_bytes()).mirrored(True, True)
qimg.convertTo(QtGui.QImage.Format.Format_RGBA8888)
data = pg.functions.ndarray_from_qimage(qimg).astype(np.float32) / 255
colors = data.reshape((-1, 4))

pg.mkQApp()
win = gl.GLViewWidget()
win.show()

rows, cols = data.shape[:2]
length = 5.0
radius = length**0.5
md = gl.MeshData.cylinder(rows=rows-1, cols=cols, radius=[radius, radius], length=length)
md.setVertexColors(colors)
item = gl.GLMeshItem(meshdata=md, smooth=False, shader='shaded')
item.translate(0, 0, -length/2)
win.addItem(item)

def update():
    win.orbit(-1, 0)
    framecnt.update()

timer = QtCore.QTimer()
timer.timeout.connect(update)
timer.start(0)

framecnt = FrameCounter()
framecnt.sigFpsUpdate.connect(lambda fps: win.setWindowTitle(f'{fps:.1f} fps'))

pg.exec()

import argparse
from PySide6 import QtGui, QtOpenGL
import numpy as np
from OpenGL import GL
from OpenGL import GLU
from OpenGL.GL import shaders

def debug_callback(source, type_, id_, severity, length, message, userParam):
    debug_enums = [x for x in dir(GL) if x.startswith("GL_DEBUG_") and not x.endswith("_KHR")]
    lut = { getattr(GL, name) : name[9:] for name in debug_enums }
        
    print(lut[source], lut[type_], hex(id_), lut[severity], message)

# these are extremely old-style shaders
VERT_SRC = """
void main() {
    gl_Position = gl_ProjectionMatrix * gl_Vertex;
    gl_FrontColor = gl_Color;
}
"""
FRAG_SRC = """
void main() {
    gl_FragColor = gl_Color;
}
"""

class GLWindow(QtOpenGL.QOpenGLWindow):
    def __init__(self, mode, useShader):
        super().__init__()

        self.mode = mode
        self.useShader = useShader

        sfmt = QtGui.QSurfaceFormat()
        sfmt.setOption(QtGui.QSurfaceFormat.FormatOption.DebugContext)
        self.setFormat(sfmt)

    def initializeGL(self):
        context = self.context()
        format = context.format()

        if format.version() >= (4, 3):
            GL.glEnable(GL.GL_DEBUG_OUTPUT)
            callback = GL.GLDEBUGPROC(debug_callback)
            GL.glDebugMessageCallback(callback, None)
            self.callback = callback    # segfault on Linux if we don't hold a reference

        triangle = np.array([
            [ -0.25, -0.25, 0], [ 0.25, -0.25, 0], [ 0.0, 0.25, 0],
        ], dtype=np.float32)

        centers = [[-0.5, -0.5, 0], [0.5, 0, 0], [-0.5, 0.5, 0]]
        self.vtx = np.vstack([triangle + center for center in centers]).astype(np.float32)
        self.colors = np.ones((3, 3), dtype=np.float32)

        if self.mode == 'client':
            GL.glVertexPointer(3, GL.GL_FLOAT, 0, self.vtx)
            GL.glEnableClientState(GL.GL_VERTEX_ARRAY)

        elif self.mode == 'server':
            self.vbo = GL.glGenBuffers(1)
            GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.vbo)
            GL.glBufferData(GL.GL_ARRAY_BUFFER, self.vtx.nbytes, self.vtx, GL.GL_STATIC_DRAW)
            GL.glVertexPointer(3, GL.GL_FLOAT, 0, None)
            GL.glEnableClientState(GL.GL_VERTEX_ARRAY)

        if self.useShader:
            self.program = shaders.compileProgram(
                shaders.compileShader(VERT_SRC, GL.GL_VERTEX_SHADER),
                shaders.compileShader(FRAG_SRC, GL.GL_FRAGMENT_SHADER),
            )
            GL.glUseProgram(self.program)

    def drawTriangle(self, index):
        GL.glColor3f(*self.colors[index])

        if self.mode in ['client', 'server']:
            GL.glDrawArrays(GL.GL_TRIANGLES, 3*index, 3)
        else:
            GL.glBegin(GL.GL_TRIANGLES)
            GL.glVertex3f(*self.vtx[3*index+0])
            GL.glVertex3f(*self.vtx[3*index+1])
            GL.glVertex3f(*self.vtx[3*index+2])
            GL.glEnd()

    def paintGL(self):
        GL.glClearColor(0, 0, 0, 1)
        GL.glClear(GL.GL_COLOR_BUFFER_BIT)

        GL.glInitNames()
        GL.glPushName(0)

        for idx in range(3):
            GL.glLoadName(idx)
            self.drawTriangle(idx)

    def mouseReleaseEvent(self, ev):
        self.makeCurrent()

        lpos = ev.position() if hasattr(ev, 'position') else ev.localPos()

        GL.glSelectBuffer(100)
        GL.glRenderMode(GL.GL_SELECT)

        vp = (0, 0, self.width(), self.height())
        GL.glMatrixMode(GL.GL_PROJECTION)
        GL.glLoadIdentity()
        GLU.gluPickMatrix(lpos.x(), vp[3] - lpos.y(), 5, 5, vp)

        self.paintGL()

        GL.glMatrixMode(GL.GL_PROJECTION)
        GL.glLoadIdentity()

        hits = GL.glRenderMode(GL.GL_RENDER)
        for hit in hits:
            print(hit.names)
            for idx in hit.names:
                self.colors[idx][idx] = 1.0 - self.colors[idx][idx]
            self.update()

parser = argparse.ArgumentParser()
parser.add_argument('--mode', choices=['immediate', 'client', 'server'], default='immediate')
parser.add_argument('--use-shader', action='store_true')
ARGS = parser.parse_args()

app = QtGui.QGuiApplication([])
win = GLWindow(ARGS.mode, ARGS.use_shader)
win.resize(400, 300)
win.show()
app.exec()