import numpy as np
from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GLU import *
from typing import Protocol

class ImageRenderer:
    def __init__(self):
        # Setup shader and quad
        self.shader = self._init_shader()
        self.vao = self._init_quad()

        # Get uniform locations from the shader
        self.uTransform = glGetUniformLocation(self.shader, "uTransform")
        self.uAlpha = glGetUniformLocation(self.shader, "uAlpha")

        # Use the shader program and set default texture unit to 0
        glUseProgram(self.shader)
        glUniform1i(glGetUniformLocation(self.shader, "uTexture"), 0)

        # Setup transition
        self.transition = None

    def set_transition(self, transition_cls):
        self.transition = transition_cls(self)

    @staticmethod
    def compile_shader(source, shader_type):
        shader = glCreateShader(shader_type)             # Create a new shader object
        glShaderSource(shader, source)                   # Attach the shader source code
        glCompileShader(shader)                          # Compile the shader
        if not glGetShaderiv(shader, GL_COMPILE_STATUS): # Check if compilation succeeded
            raise RuntimeError(glGetShaderInfoLog(shader).decode()) # Raise error with log if failed
        return shader

    @staticmethod
    def _init_shader():
        vertex_src = """
        #version 330 core
        layout (location = 0) in vec2 aPos;      // Vertex position
        layout (location = 1) in vec2 aTexCoord; // Texture coordinates
        uniform mat4 uTransform;                 // Transformation matrix
        out vec2 TexCoord;                       // Output texture coordinate

        void main() {
            gl_Position = uTransform * vec4(aPos, 0.0, 1.0); // Apply transformation
            TexCoord = aTexCoord;                            // Pass tex coord to fragment shader
        }
        """
        fragment_src = """
        #version 330 core
        in vec2 TexCoord;           // Interpolated texture coordinates
        out vec4 FragColor;         // Final fragment color
        uniform sampler2D uTexture; // Texture sampler
        uniform float uAlpha;       // Global alpha for transparency

        void main() {
            vec4 texColor = texture(uTexture, TexCoord);         // Sample texture
            FragColor = vec4(texColor.rgb, texColor.a * uAlpha); // Apply alpha blending
        }
        """

        # Compile and link shaders into a program
        vs = ImageRenderer.compile_shader(vertex_src, GL_VERTEX_SHADER)
        fs = ImageRenderer.compile_shader(fragment_src, GL_FRAGMENT_SHADER)
        prog = glCreateProgram()
        glAttachShader(prog, vs)
        glAttachShader(prog, fs)
        glLinkProgram(prog)
        return prog

    @staticmethod
    def _init_quad():
        # Define a full-screen quad with positions and texture coordinates
        quad = np.array([
            -1, -1,  0, 0,  # Bottom-left
             1, -1,  1, 0,  # Bottom-right
            -1,  1,  0, 1,  # Top-left
             1,  1,  1, 1,  # Top-right
        ], dtype=np.float32)

        # Create and bind a Vertex Array Object
        vao = glGenVertexArrays(1)
        vbo = glGenBuffers(1)
        glBindVertexArray(vao)
        glBindBuffer(GL_ARRAY_BUFFER, vbo)
        glBufferData(GL_ARRAY_BUFFER, quad.nbytes, quad, GL_STATIC_DRAW)

        # Setup vertex attributes: position (location = 0)
        glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 16, ctypes.c_void_p(0))
        glEnableVertexAttribArray(0)

        # Setup vertex attributes: texture coordinates (location = 1)
        glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 16, ctypes.c_void_p(8))
        glEnableVertexAttribArray(1)
        return vao

    def draw_image(self, tex, win_w, win_h, alpha):
        glUseProgram(self.shader)
        glBindVertexArray(self.vao)
        glBindTexture(GL_TEXTURE_2D, tex.id)

        # Calculate aspect ratios
        img_aspect = tex.width / tex.height
        win_aspect = win_w / win_h

        # Calculate scaling factors to preserve image aspect ratio
        if img_aspect > win_aspect:
            sx = 1.0
            sy = win_aspect / img_aspect
        else:
            sx = img_aspect / win_aspect
            sy = 1.0

        # Create transformation matrix for aspect-ratio-correct rendering
        transform = np.array([
            [sx, 0, 0, 0],
            [0, sy, 0, 0],
            [0,  0, 1, 0],
            [0,  0, 0, 1],
        ], dtype=np.float32)

        # Pass transformation and alpha to the shader
        glUniformMatrix4fv(self.uTransform, 1, GL_FALSE, transform)
        glUniform1f(self.uAlpha, alpha)

        # Draw the textured quad
        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4)

    def draw_static(self, tex, win_w, win_h, alpha):
        # Set the background color to black
        glClearColor(0.0, 0.0, 0.0, 1.0)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

        self.draw_image(tex, win_w, win_h, alpha)

        glutSwapBuffers()

    def draw_transition(self, tex_start, tex_end, win_w, win_h, delta_time, transition_time, transition_duration, reversed):
        assert self.transition, "No transition has been set"

        # Set the background color to black
        glClearColor(0.0, 0.0, 0.0, 1.0)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

        self.transition.draw(tex_start, tex_end, win_w, win_h, delta_time, transition_time, transition_duration, reversed)

        glutSwapBuffers()


class Transition(Protocol):
    def __init__(self, renderer):
        self.renderer = renderer

    def draw(self, tex_start, tex_end, win_w, win_h, delta_time, transition_time, transition_duration, reversed):
        pass


class TransitionMix(Transition):
    def draw(self, tex_start, tex_end, win_w, win_h, delta_time, transition_time, transition_duration, reversed):
        # Update alpha value for fade effect
        alpha = transition_time / transition_duration
        if alpha > 1.0:
            alpha = 1.0

        # Draw the images on top of one another
        self.renderer.draw_image(tex_start, win_w, win_h, 1 - alpha)  # TODO instead of decreasing alpha, draw transparent letterboxes
        self.renderer.draw_image(tex_end,   win_w, win_h, alpha)