'''
    Copyright (C) 2020 - 2025 Akaneyu

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
'''

import gpu
from gpu.shader import create_from_info
from gpu_extras.batch import batch_for_shader
from mathutils import Matrix
import numpy as np

default_vertex_shader = '''
void main()
{
    gl_Position = ModelViewProjectionMatrix * vec4(pos, 0, 1.0);
}
'''

default_fragment_shader = '''
void main()
{
    fragColor = color;
}
'''

dotted_line_vertex_shader = '''
void main()
{
    arcLengthInter = arcLength;

    gl_Position = ModelViewProjectionMatrix * vec4(pos, 0, 1.0);
}
'''

dotted_line_fragment_shader = '''
void main()
{
    if (step(sin((arcLengthInter + offset) * scale), 0.5) == 1) {
        fragColor = color1;
    } else {
        fragColor = color2;
    }
}
'''

class UIRenderer:
    def __init__(self):
        default_shader_info = gpu.types.GPUShaderCreateInfo()
        default_shader_info.push_constant('MAT4', 'ModelViewProjectionMatrix')
        default_shader_info.push_constant('VEC4', 'color')
        default_shader_info.vertex_in(0, 'VEC2', 'pos')
        default_shader_info.fragment_out(0, 'VEC4', 'fragColor')
        default_shader_info.vertex_source(default_vertex_shader)
        default_shader_info.fragment_source(default_fragment_shader)
        self.default_shader = create_from_info(default_shader_info)
        self.default_shader_u_color = self.default_shader.uniform_from_name('color')

        dotted_line_shader_inter = gpu.types.GPUStageInterfaceInfo("dotted_line")
        dotted_line_shader_inter.smooth('FLOAT', "arcLengthInter")

        dotted_line_shader_info = gpu.types.GPUShaderCreateInfo()
        dotted_line_shader_info.push_constant('MAT4', 'ModelViewProjectionMatrix')
        dotted_line_shader_info.push_constant('FLOAT', 'scale')
        dotted_line_shader_info.push_constant('FLOAT', 'offset')
        dotted_line_shader_info.push_constant('VEC4', 'color1')
        dotted_line_shader_info.push_constant('VEC4', 'color2')
        dotted_line_shader_info.vertex_in(0, 'VEC2', 'pos')
        dotted_line_shader_info.vertex_in(1, 'FLOAT', 'arcLength')
        dotted_line_shader_info.vertex_out(dotted_line_shader_inter)
        dotted_line_shader_info.fragment_out(0, 'VEC4', 'fragColor')
        dotted_line_shader_info.vertex_source(dotted_line_vertex_shader)
        dotted_line_shader_info.fragment_source(dotted_line_fragment_shader)
        self.dotted_line_shader = create_from_info(dotted_line_shader_info)
        self.dotted_line_shader_u_color1 = self.dotted_line_shader.uniform_from_name("color1")
        self.dotted_line_shader_u_color2 = self.dotted_line_shader.uniform_from_name("color2")

    def render_border(self, pos1, pos2):
        batch = batch_for_shader(self.default_shader, 'LINES',
            {"pos": [pos1, pos2]})

        self.default_shader.bind()

        self.default_shader.uniform_vector_float(self.default_shader_u_color,
                np.array([1.0, 0.0, 1.0, 1.0], 'f'), 4)

        batch.draw(self.default_shader)

    def render_arrow(self, center, angle):
        prev_blend = gpu.state.blend_get()
        gpu.state.blend_set('ALPHA')

        with gpu.matrix.push_pop():
            gpu.matrix.translate(center)
            gpu.matrix.multiply_matrix(
                    Matrix.Rotation(angle, 4, 'Z'))

            verts = [
                (0, -50),
                (100, -50),
                (0, 50),
                (100, 50),
                (100, 0),
                (200, 0),
                (100, 100),
                (100, -100)
            ]

            indices = [
                (0, 1, 2),
                (2, 1, 3),
                (4, 5, 6),
                (4, 5, 7)
            ]

            batch = batch_for_shader(self.default_shader, 'TRIS',
                {"pos": verts}, indices=indices)

            self.default_shader.bind()

            self.default_shader.uniform_vector_float(self.default_shader_u_color,
                    np.array([1.0, 1.0, 1.0, 0.5], 'f'), 4)

            batch.draw(self.default_shader)

        gpu.state.blend_set(prev_blend)

    def render_brush_frame(self, pos, radius):
        prev_line_width = gpu.state.line_width_get()
        gpu.state.line_width_set(2.0)

        verts = self.create_brush_frame_vertices(pos, radius)

        arc_lengths = [0]
        for a, b in zip(verts[:-1], verts[1:]):
            arc_lengths.append(arc_lengths[-1] + np.linalg.norm(a - b))
    
        batch = batch_for_shader(self.dotted_line_shader, 'LINE_STRIP',
            {"pos": verts, "arcLength": arc_lengths})

        self.dotted_line_shader.bind()

        self.dotted_line_shader.uniform_float("scale", 0.6)
        self.dotted_line_shader.uniform_float("offset", 0)
        self.dotted_line_shader.uniform_vector_float(self.dotted_line_shader_u_color1,
            np.array([1.0, 1.0, 1.0, 0.5], 'f'), 4)
        self.dotted_line_shader.uniform_vector_float(self.dotted_line_shader_u_color2,
            np.array([0.0, 0.0, 0.0, 0.5], 'f'), 4)

        batch.draw(self.dotted_line_shader)

        gpu.state.line_width_set(prev_line_width)

    def create_brush_frame_vertices(self, pos, radius):
        segs = 32

        theta = 2.0 * np.pi / segs
        c = np.cos(theta)
        s = np.sin(theta)

        x = radius 
        y = 0

        verts = []
        for i in range(segs):
            verts.append((x + pos[0], y + pos[1])) 
            
            t = x
            x = c * x - s * y
            y = s * t + c * y

        verts.append(verts[0])

        return np.array(verts, 'f')

    def render_info_box(self, pos1, pos2):
        prev_blend = gpu.state.blend_get()
        gpu.state.blend_set('ALPHA')

        verts = [
            pos1,
            (pos2[0], pos1[1]),
            (pos1[0], pos2[1]),
            pos2
        ]

        indices = [
            (0, 1, 2),
            (2, 1, 3)
        ]

        batch = batch_for_shader(self.default_shader, 'TRIS',
            {"pos": verts}, indices=indices)

        self.default_shader.bind()

        self.default_shader.uniform_vector_float(self.default_shader_u_color,
                np.array([0, 0, 0, 0.7], 'f'), 4)

        batch.draw(self.default_shader)

        gpu.state.blend_set(prev_blend)
