from collections.abc import Iterable
from enum import Enum
from typing import Literal, cast, overload

import gi
import OpenGL.GL as GL
from fabric import Property, Signal
from fabric.widgets.widget import Widget
from OpenGL.GL.shaders import compileProgram, compileShader

gi.require_version("Gtk", "3.0")
from gi.repository import Gdk, GdkPixbuf, GLib, Gtk


class ShadertoyUniformType(Enum):
    FLOAT = 1
    INTEGER = 2
    VECTOR = 3
    TEXTURE = 4

class ShadertoyCompileError(Exception): ...

class Shadertoy(Gtk.GLArea, Widget):
    @Signal
    def ready(self) -> None: ...

    @Property(str, "read-write")
    def shader_buffer(self) -> str:
        return self._shader_buffer

    @shader_buffer.setter
    def shader_buffer(self, shader_buffer: str) -> None:
        self._shader_buffer = shader_buffer
        if not self._ready:
            return
        self._shader_uniforms.clear()
        self.do_realize()
        self.queue_draw()
        return

    DEFAULT_VERTEX_SHADER = """

    in vec2 position;

    void main() {
        gl_Position = vec4(position, 0.0, 1.0);
    }
    """

    DEFAULT_FRAGMENT_UNIFORMS = """

    uniform vec3 iResolution;           // viewport resolution (in pixels)
    uniform float iTime;                 // shader playback time (in seconds)
    uniform float iTimeDelta;            // render time (in seconds)
    uniform float iFrameRate;            // shader frame rate
    uniform int iFrame;                  // shader playback frame
    uniform float iChannelTime[4];       // channel playback time (in seconds)
    uniform vec3 iChannelResolution[4];  // channel resolution (in pixels)
    uniform vec4 iMouse;                 // mouse pixel coords. xy: current (if MLB down), zw: click
    uniform sampler2D iChannel0;         // input channel. XX = 2D/Cube
    uniform sampler2D iChannel1;
    uniform sampler2D iChannel2;
    uniform sampler2D iChannel3;
    uniform vec4 iDate;                  // (year, month, day, time in seconds)
    uniform float iSampleRate;           // sound sample rate (i.e., 44100)

    """

    FRAGMENT_MAIN_FUNCTION = """
    void main() {
        mainImage(gl_FragColor, gl_FragCoord.xy);
    }
    """

    def __init__(
        self,
        shader_buffer: str,
        shader_uniforms: list[
            tuple[
                str,
                ShadertoyUniformType,
                bool | float | int | tuple[float, ...] | GdkPixbuf.Pixbuf,
            ]
        ]
        | None = None,
        name: str | None = None,
        visible: bool = True,
        all_visible: bool = False,
        style: str | None = None,
        style_classes: Iterable[str] | str | None = None,
        tooltip_text: str | None = None,
        tooltip_markup: str | None = None,
        h_align: Literal["fill", "start", "end", "center", "baseline"]
        | Gtk.Align
        | None = None,
        v_align: Literal["fill", "start", "end", "center", "baseline"]
        | Gtk.Align
        | None = None,
        h_expand: bool = False,
        v_expand: bool = False,
        size: Iterable[int] | int | None = None,
        **kwargs,
    ):
        Gtk.GLArea.__init__(
            self
        )
        Widget.__init__(
            self,
            name,
            visible,
            all_visible,
            style,
            style_classes,
            tooltip_text,
            tooltip_markup,
            h_align,
            v_align,
            h_expand,
            v_expand,
            size,
            **kwargs,
        )
        self._shader_buffer = shader_buffer
        self._shader_uniforms = shader_uniforms or []

        self.set_required_version(3, 3)
        self.set_has_depth_buffer(False)
        self.set_has_stencil_buffer(False)

        self._ready = False
        self._program = None
        self._vao = None
        self._quad_vbo = None
        self._texture_units = {}

        self._start_time = GLib.get_monotonic_time() / 1e6
        self._frame_time = self._start_time
        self._frame_count = 0

        self._tick_id = self.add_tick_callback(lambda *_: (self.queue_draw(), True)[1])

    def do_bake_program(self):
        try:
            vertex_shader = compileShader(
                self.DEFAULT_VERTEX_SHADER, GL.GL_VERTEX_SHADER
            )
            fragment_shader = compileShader(
                self.DEFAULT_FRAGMENT_UNIFORMS
                + self._shader_buffer
                + self.FRAGMENT_MAIN_FUNCTION,
                GL.GL_FRAGMENT_SHADER,
            )
        except Exception as e:
            raise ShadertoyCompileError(
                f"couldn't compile the provided shader, OpenGL error:\n {e}"
            )

        return compileProgram(vertex_shader, fragment_shader)

    def do_realize(self, *_):
        Gtk.GLArea.do_realize(self)
        if not self._ready:
            ctx = self.get_context()
            if (err := self.get_error()) or not ctx:
                raise RuntimeError(
                    f"couldn't initialize the drawing context, error: {err or 'context is None'}"
                )

            ctx.make_current()

        if self._program:
            GL.glDeleteProgram(self._program)
            self._program = None
        self._program = self.do_bake_program()

        GL.glEnable(GL.GL_BLEND)
        GL.glBlendFunc(GL.GL_SRC_ALPHA, GL.GL_ONE_MINUS_SRC_ALPHA)

        self._quad_vbo = GL.glGenBuffers(1)
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self._quad_vbo)

        quad_verts = (-1.0, -1.0, 1.0, -1.0, -1.0, 1.0, 1.0, 1.0)
        array_type = GL.GLfloat * len(quad_verts)

        GL.glBufferData(
            GL.GL_ARRAY_BUFFER,
            len(quad_verts) * 4,
            array_type(*quad_verts),
            GL.GL_STATIC_DRAW,
        )

        self._vao = GL.glGenVertexArrays(1)
        GL.glBindVertexArray(self._vao)

        position = GL.glGetAttribLocation(self._program, "position")
        GL.glEnableVertexAttribArray(position)
        GL.glVertexAttribPointer(position, 2, GL.GL_FLOAT, GL.GL_FALSE, 0, None)

        for uname, utype, uvalue in self._shader_uniforms:
            self.set_uniform(uname, utype, uvalue)

        self._ready = True
        self.ready()
        return

    def do_get_timing(self) -> tuple[float, float, float]:
        current_time = GLib.get_monotonic_time() / 1e6
        delta_time = current_time - self._frame_time
        return current_time, delta_time, (1.0 / delta_time) if delta_time > 0 else 0.0

    def do_post_render(self, time: float):
        self._frame_time = time
        self._frame_count += 1
        return

    def do_render(self, ctx: Gdk.GLContext):
        if not self._program:
            if self._tick_id:
                self.remove_tick_callback(self._tick_id)
            self._tick_id = 0
            return False

        GL.glUseProgram(self._program)

        GL.glClear(GL.GL_COLOR_BUFFER_BIT)

        alloc = self.get_allocation()
        width: int = alloc.width
        height: int = alloc.height
        mouse_pos = cast(tuple[int, int], self.get_pointer())

        current_time, delta_time, frame_rate = self.do_get_timing()

        self.set_uniform(
            "iTime", ShadertoyUniformType.FLOAT, current_time - self._start_time
        )
        self.set_uniform("iFrame", ShadertoyUniformType.INTEGER, self._frame_count)
        self.set_uniform("iTimeDelta", ShadertoyUniformType.FLOAT, delta_time)
        self.set_uniform("iFrameRate", ShadertoyUniformType.FLOAT, frame_rate)
        self.set_uniform(
            "iResolution", ShadertoyUniformType.VECTOR, (width, height, 1.0)
        )
        self.set_uniform(
            "iMouse",
            ShadertoyUniformType.VECTOR,
            (mouse_pos[0], height - mouse_pos[1], 0, 0),
        )

        GL.glBindVertexArray(self._vao)
        GL.glDrawArrays(GL.GL_TRIANGLE_STRIP, 0, 4)
        self.do_post_render(current_time)
        return True

    def do_resize(self, width: int, height: int):
        Gtk.GLArea.do_resize(self, width, height)
        GL.glViewport(0, 0, width, height)
        return

    @overload
    def set_uniform(
        self, name: str, type: Literal[ShadertoyUniformType.FLOAT], value: float
    ): ...

    @overload
    def set_uniform(
        self, name: str, type: Literal[ShadertoyUniformType.INTEGER], value: int
    ): ...

    @overload
    def set_uniform(
        self,
        name: str,
        type: Literal[ShadertoyUniformType.VECTOR],
        value: tuple[float, ...],
    ): ...

    @overload
    def set_uniform(
        self,
        name: str,
        type: Literal[ShadertoyUniformType.TEXTURE],
        value: GdkPixbuf.Pixbuf,
    ): ...

    def set_uniform(
        self,
        name: str,
        type: ShadertoyUniformType,
        value: bool | float | int | tuple[float, ...] | GdkPixbuf.Pixbuf,
    ):
        if not self._program:
            raise RuntimeError("the shader program is not initialized")
        GL.glUseProgram(self._program)
        location = GL.glGetUniformLocation(self._program, name)
        match type:
            case ShadertoyUniformType.VECTOR:
                value = cast(tuple[float, ...], value)
                (
                    GL.glUniform2f
                    if (vlen := len(value)) == 2
                    else GL.glUniform3f
                    if vlen == 3
                    else GL.glUniform4f
                )(location, *value)
            case ShadertoyUniformType.FLOAT:
                GL.glUniform1f(location, value)
            case ShadertoyUniformType.INTEGER:
                GL.glUniform1i(location, value)
            case ShadertoyUniformType.TEXTURE:

                value = cast(GdkPixbuf.Pixbuf, value).flip(False)
                format = GL.GL_RGBA if value.get_has_alpha() else GL.GL_RGB

                if name not in self._texture_units:
                    texture = GL.glGenTextures(1)
                    self._texture_units[name] = (len(self._texture_units), texture)
                else:
                    texture_unit, texture = self._texture_units[name]

                texture_unit = self._texture_units[name][0]
                GL.glActiveTexture(GL.GL_TEXTURE0 + texture_unit)
                GL.glBindTexture(GL.GL_TEXTURE_2D, texture)

                GL.glTexParameteri(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_WRAP_S, GL.GL_REPEAT)
                GL.glTexParameteri(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_WRAP_T, GL.GL_REPEAT)
                GL.glTexParameteri(
                    GL.GL_TEXTURE_2D, GL.GL_TEXTURE_MIN_FILTER, GL.GL_LINEAR
                )
                GL.glTexParameteri(
                    GL.GL_TEXTURE_2D, GL.GL_TEXTURE_MAG_FILTER, GL.GL_LINEAR
                )

                GL.glTexImage2D(
                    GL.GL_TEXTURE_2D,
                    0,
                    format,
                    value.get_width(),
                    value.get_height(),
                    0,
                    format,
                    GL.GL_UNSIGNED_BYTE,
                    value.get_pixels(),
                )
                GL.glGenerateMipmap(GL.GL_TEXTURE_2D)

                GL.glUniform1i(location, texture_unit)