import pygame as pg
import numpy as np

import pyassimp
import glm

from OpenGL.GL import *

vertex_shader = """
#version 330 core

layout (location = 0) in vec3 aPos;

uniform mat4 model;
uniform mat4 projection;
uniform mat4 view;

void main(){
    gl_Position = projection * view * model * vec4(aPos, 1.0);
}

"""

fragment_shader = """
#version 330 core

out vec4 fragColor;

void main(){
    fragColor = vec4(1.0, 0.0, 1.0, 1.0);
}

"""

class Camera3d():
    def __init__(self, position=glm.vec3(0.0), rotation=glm.vec3(0.0)):
        self.position = position
        self.rotation = rotation
        self.projection_matrix = glm.perspective(glm.radians(75), 800 / 600, 0.1, 1024)

    @property
    def forward(self):
        forward = glm.vec3()
        forward.x = glm.cos(glm.radians(self.rotation.y)) * glm.cos(glm.radians(self.rotation.x))
        forward.y = glm.sin(glm.radians(self.rotation.x))
        forward.z = glm.sin(glm.radians(self.rotation.y)) * glm.cos(glm.radians(self.rotation.x))
        return glm.normalize(forward)

    @property
    def right(self):
        return glm.normalize(glm.cross(self.forward, glm.vec3(0, 1, 0)))

    @property
    def up(self):
        return glm.normalize(glm.cross(self.right, self.forward))

    def update(self):
        keys = pg.key.get_pressed()

        if keys[pg.K_w]:
            self.position += glm.vec3(0, 0, 1)
        elif keys[pg.K_s]:
            self.position += glm.vec3(0, 0, -1)
        if keys[pg.K_a]:
            self.position += glm.vec3(-1, 0, 0)
        elif keys[pg.K_d]:
            self.position += glm.vec3(1, 0, 0)

    def get_projection_matrix(self):
        return self.projection_matrix

    def get_view_matrix(self):
        return glm.lookAt(self.position, self.position + self.forward, self.up)

class Mesh3d():
    def __init__(self):
        self.position = glm.vec3(0.0)
        self.rotation = glm.vec3(0.0)

        self.vertices = np.array([], dtype=np.float32)
        self.texcoords = np.array([], dtype=np.uint32)
        self.normals = np.array([], dtype=np.float32)
        self.faces = np.array([], dtype=np.uint32)

        self.vao, self.vbo, self.ebo = 0, 0, 0
        self.shader = self.create_shader_program(vertex_shader, fragment_shader)

    def compile_shader(self, source: str, shader_type: GL_SHADER_TYPE):
        shader = glCreateShader(shader_type)
        glShaderSource(shader, source)
        glCompileShader(shader)
        return shader

    def create_shader_program(self, vertex_source: str, fragment_source: str):
        vertex_shader = self.compile_shader(vertex_source, GL_VERTEX_SHADER)
        fragment_shader = self.compile_shader(fragment_source, GL_FRAGMENT_SHADER)

        shader_program = glCreateProgram()
        glAttachShader(shader_program, vertex_shader)
        glAttachShader(shader_program, fragment_shader)
        glLinkProgram(shader_program)

        glDeleteShader(vertex_shader)
        glDeleteShader(fragment_shader)

        return shader_program

    def get_model_matrix(self):
        model = glm.mat4(1.0)
        model = glm.translate(model, self.position)
        return model

    def setup_buffers(self):
        self.vao = glGenVertexArrays(1)
        glBindVertexArray(self.vao)

        self.vbo = glGenBuffers(1)
        glBindBuffer(GL_ARRAY_BUFFER, self.vbo)
        glBufferData(GL_ARRAY_BUFFER, self.vertices.nbytes,
                     self.vertices, GL_STATIC_DRAW)

        self.ebo = glGenBuffers(1)
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.ebo)
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.faces.nbytes,
                     self.faces, GL_STATIC_DRAW)

        glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE,
                              3 * self.vertices.dtype.itemsize, None)
        glEnableVertexAttribArray(0)

        glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE,
                              2 * self.faces.dtype.itemsize, None)
        glEnableVertexAttribArray(1)

        glBindVertexArray(0)

    def render(self, camera: Camera3d):
        if not self.vao:
            return
        glUseProgram(self.shader)

        model_loc = glGetUniformLocation(self.shader, "model")
        glUniformMatrix4fv(model_loc, 1, GL_FALSE,
                           glm.value_ptr(self.get_model_matrix()))

        projection_loc = glGetUniformLocation(self.shader, "projection")
        glUniformMatrix4fv(projection_loc, 1, GL_FALSE,
                           glm.value_ptr(camera.get_projection_matrix()))

        view_loc = glGetUniformLocation(self.shader, "view")
        glUniformMatrix4fv(view_loc, 1, GL_FALSE,
                           glm.value_ptr(camera.get_view_matrix()))

        glBindVertexArray(self.vao)

        glDrawElements(GL_TRIANGLES, len(self.faces) * 3, GL_UNSIGNED_INT, None)

        glBindVertexArray(0)
        glUseProgram(0)

    def load_from(self, filename: str):
        with pyassimp.load(filename) as scene:
            mesh = scene.meshes[0]
            self.vertices = mesh.vertices
            self.texcoords = mesh.texturecoords
            self.normals = mesh.normals
            self.faces = mesh.faces
        self.setup_buffers()
        return self

class App():
    def __init__(self):
        self.win = pg.display.set_mode((800, 600), pg.OPENGL | pg.DOUBLEBUF, vsync=True)
        self.clock = pg.time.Clock()

        self.mesh = Mesh3d().load_from('res/teapot.obj')
        self.camera = Camera3d(glm.vec3(0, 0, -3))

    def __events(self):
        for e in pg.event.get():
            if e.type == pg.QUIT:
                self.__is_running = False

    def __update(self):
        self.camera.update()

    def __render(self):
        glClearColor(0.05, 0.05, 0.05, 1)
        glClear(GL_COLOR_BUFFER_BIT)

        self.mesh.render(self.camera)

        pg.display.flip()

    def run(self):
        self.__is_running = True

        while self.__is_running:
            self.__events()
            self.__update()
            self.__render()
            self.clock.tick(0)

        pg.quit()

if __name__ == '__main__':
    App().run()

这段代码应该只绘制从文件加载的网格，但是当我运行它时，我只看到空。