# Copyright (C) 2022 RozK
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import ctypes
from array import array
from pathlib import Path

_engine = ctypes.cdll.LoadLibrary(Path(__file__).parent / "engine.so")

# types.hpp

def _flag(x):
    return 1 << x

def buffer(type, size):
    return (type * size)()

def _void_addr(x):
    return x.buffer_info()[0]

def _ubyte_addr(x):
    assert x.typecode == 'B'
    return x.buffer_info()[0]

def _ushort_addr(x):
    assert x.typecode == 'H'
    return x.buffer_info()[0]

def _float_addr(x):
    assert x.typecode == 'f'
    return x.buffer_info()[0]

# math.hpp

class vec3(ctypes.Structure):
    _fields_ = ('x', ctypes.c_float), ('y', ctypes.c_float), ('z', ctypes.c_float)

    def __iter__(self):
        yield self.x
        yield self.y
        yield self.z

    def set(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z

class vec4(ctypes.Structure):
    _fields_ = ('x', ctypes.c_float), ('y', ctypes.c_float), ('z', ctypes.c_float), ('w', ctypes.c_float)

    def __iter__(self):
        yield self.x
        yield self.y
        yield self.z
        yield self.w

    def set(self, x, y, z, w):
        self.x = x
        self.y = y
        self.z = z
        self.w = w

class mat3(ctypes.Structure):
    _fields_ = ('right', vec3), ('forward', vec3), ('up', vec3)

    def __iter__(self):
        yield self.right
        yield self.forward
        yield self.up

    def set(self, right, forward, up):
        self.right = right
        self.forward = forward
        self.up = up

class mat4(ctypes.Structure):
    _fields_ = ('right', vec4), ('forward', vec4), ('up', vec4), ('origin', vec4)

    def __iter__(self):
        yield self.right
        yield self.forward
        yield self.up
        yield self.origin

    def set(self, right, forward, up, origin):
        self.right = right
        self.forward = forward
        self.up = up
        self.origin = origin

    def set_mat3_vec3(self, axes, origin):
        self.right.set(*axes.right, 0.0)
        self.forward.set(*axes.forward, 0.0)
        self.up.set(*axes.up, 0.0)
        self.origin.set(*origin, 1.0)

    def set_axes(self, right, forward, up):
        self.right = right
        self.forward = forward
        self.up = up

    def set_axes_mat3(self, axes):
        self.right.set(*axes.right, 0.0)
        self.forward.set(*axes.forward, 0.0)
        self.up.set(*axes.up, 0.0)

    def set_origin(self, origin):
        self.origin = origin

    def set_origin_vec3(self, origin):
        self.origin.set(*origin, 1.0)

assert ctypes.sizeof(vec3) == ctypes.sizeof(ctypes.c_float) * 3
assert ctypes.sizeof(vec4) == ctypes.sizeof(ctypes.c_float) * 4
assert ctypes.sizeof(mat3) == ctypes.sizeof(vec3) * 3
assert ctypes.sizeof(mat4) == ctypes.sizeof(vec4) * 4

_vec3_p = ctypes.POINTER(vec3)
_vec4_p = ctypes.POINTER(vec4)
_mat3_p = ctypes.POINTER(mat3)
_mat4_p = ctypes.POINTER(mat4)

vec3_right   = vec3(1.0, 0.0, 0.0)
vec3_forward = vec3(0.0, 1.0, 0.0)
vec3_up      = vec3(0.0, 0.0, 1.0)
vec3_origin  = vec3(0.0, 0.0, 0.0)

vec4_right   = vec4(1.0, 0.0, 0.0, 0.0)
vec4_forward = vec4(0.0, 1.0, 0.0, 0.0)
vec4_up      = vec4(0.0, 0.0, 1.0, 0.0)
vec4_origin  = vec4(0.0, 0.0, 0.0, 1.0)

mat3_identity = mat3(
    vec3_right,
    vec3_forward,
    vec3_up)

mat4_identity = mat4(
    vec4_right,
    vec4_forward,
    vec4_up,
    vec4_origin)

vec3_rotate = _engine.rk_vec3_rotate
vec3_rotate.argtypes = (
    _vec3_p, # ret
    _vec3_p, # vec3
    _vec3_p, # axis
    ctypes.c_float) # angle

vec3_mul_vec3 = _engine.rk_vec3_mul_vec3
vec3_mul_vec3.argtypes = (
    _vec3_p, # ret
    _vec3_p, # a
    _vec3_p) # b

mat3_rotation = _engine.rk_mat3_rotation
mat3_rotation.argtypes = (
    _mat3_p, # ret
    _vec3_p, # axis
    ctypes.c_float) # angle

mat3_mul_vec3 = _engine.rk_mat3_mul_vec3
mat3_mul_vec3.argtypes = (
    _vec3_p, # ret
    _mat3_p, # a
    _vec3_p) # b

mat3_mul_mat3 = _engine.rk_mat3_mul_mat3
mat3_mul_mat3.argtypes = (
    _mat3_p, # ret
    _mat3_p, # a
    _mat3_p) # b

mat4_projection = _engine.rk_mat4_projection
mat4_projection.argtypes = (
    _mat4_p, # ret
    ctypes.c_float, # hfov
    ctypes.c_float, # ratio
    ctypes.c_float, # near
    ctypes.c_float) # far

mat4_lookat = _engine.rk_mat4_lookat
mat4_lookat.argtypes = (
    _mat4_p, # ret
    _vec3_p, # position
    _vec3_p) # lookat

mat4_mul_vec3 = _engine.rk_mat4_mul_vec3
mat4_mul_vec3.argtypes = (
    _vec3_p, # ret
    _mat4_p, # a
    _vec3_p, # b
    ctypes.c_float) # w

mat4_mul_vec4 = _engine.rk_mat4_mul_vec4
mat4_mul_vec4.argtypes = (
    _vec4_p, # ret
    _mat4_p, # a
    _vec4_p) # b

mat4_mul_mat4 = _engine.rk_mat4_mul_mat4
mat4_mul_mat4.argtypes = (
    _mat4_p, # ret
    _mat4_p, # a
    _mat4_p) # b

# display.hpp

create_display = _engine.rk_create_display
create_display.restype = ctypes.c_void_p
create_display.argtypes = (
    ctypes.c_char_p, # name
    ctypes.c_uint, # width
    ctypes.c_uint) # height

destroy_display = _engine.rk_destroy_display
destroy_display.argtypes = (
    ctypes.c_void_p,) # display

swap_buffers = _engine.rk_swap_buffers
swap_buffers.argtypes = (
    ctypes.c_void_p,) # display

# render.hpp

TEXTURE_FORMAT_SRGB8_A8 = 0
TEXTURE_FORMAT_RGBA8 = 1
TEXTURE_FORMAT_RGB10_A2 = 2
TEXTURE_FORMAT_FLOAT_32 = 3

TEXTURE_FORMAT_TYPECODE = ('B', 'B', 'I', 'f')
TEXTURE_FORMAT_NELEMS = (4, 4, 1, 1)

TEXTURE_FLAG_3D = _flag(0)
TEXTURE_FLAG_MIPMAPS = _flag(1)
TEXTURE_FLAG_MIN_NEAREST = 0
TEXTURE_FLAG_MIN_LINEAR = _flag(2)
TEXTURE_FLAG_MAG_NEAREST = 0
TEXTURE_FLAG_MAG_LINEAR = _flag(3)

VERTEX_FORMAT_VEC3_FLOAT = 1
VERTEX_FORMAT_VEC3_INT10 = 2
VERTEX_FORMAT_VEC3_UINT10 = 3
VERTEX_FORMAT_NORMALIZE = _flag(7)

def vertex_format(*format):
    return array('B', format).tobytes()

PARAM_FORMAT_VEC3_FLOAT = 1
PARAM_FORMAT_VEC3_SHORT = 2
PARAM_FORMAT_VEC3_INT10 = 3
PARAM_FORMAT_MAT3_FLOAT = 4
PARAM_FORMAT_MAT3_INT10 = 5
PARAM_FORMAT_NORMALIZE = _flag(7)

def params_format(*format):
    return array('B', format).tobytes()

INSTANCE_FLAG_SPAWNED = _flag(0)
INSTANCE_FLAG_VISIBLE = _flag(1)

BATCH_MAX_SIZE = 65536

render_initialize = _engine.rk_render_initialize
render_terminate = _engine.rk_render_terminate

_load_shader = _engine.rk_load_shader
_load_shader.restype = ctypes.c_void_p
_load_shader.argtypes = (
    ctypes.c_uint, # vert_nlines
    ctypes.POINTER(ctypes.c_char_p), # vert_lines
    ctypes.c_uint, # frag_nlines
    ctypes.POINTER(ctypes.c_char_p)) # frag_lines

def load_shader(vert_lines, frag_lines):
    vert_nlines = len(vert_lines)
    vert_type = ctypes.c_char_p * vert_nlines
    vert_lines = vert_type(*map(ctypes.c_char_p, vert_lines))
    frag_nlines = len(frag_lines)
    frag_type = ctypes.c_char_p * frag_nlines
    frag_lines = frag_type(*map(ctypes.c_char_p, frag_lines))
    return _load_shader(vert_nlines, vert_lines, frag_nlines, frag_lines)

resolve_input = _engine.rk_resolve_input
resolve_input.restype = ctypes.c_void_p
resolve_input.argtypes = (
    ctypes.c_void_p, # shader
    ctypes.c_char_p) # name

resolve_param = _engine.rk_resolve_param
resolve_param.restype = ctypes.c_void_p
resolve_param.argtypes = (
    ctypes.c_void_p, # shader
    ctypes.c_char_p) # name

_create_texture = _engine.rk_create_texture
_create_texture.restype = ctypes.c_void_p
_create_texture.argtypes = (
    ctypes.c_uint, # format
    ctypes.c_uint, # width
    ctypes.c_uint, # height
    ctypes.c_uint, # nlevels
    ctypes.c_uint, # flags
    ctypes.c_void_p) # pixels

def create_texture(format, width, height, nlevels, flags, pixels):
    assert pixels.typecode == TEXTURE_FORMAT_TYPECODE[format]
    assert len(pixels) == width * height * max(1, nlevels) * TEXTURE_FORMAT_NELEMS[format]
    return _create_texture(format, width, height, nlevels, flags, _void_addr(pixels))

_create_triangles = _engine.rk_create_triangles
_create_triangles.restype = ctypes.c_void_p
_create_triangles.argtypes = (
    ctypes.c_uint, # nvertices
    ctypes.c_void_p) # vertices

def create_triangles(vertices):
    assert len(vertices) % 9 == 0
    return _create_triangles(len(vertices) // 3, _float_addr(vertices))

_create_vertices = _engine.rk_create_vertices
_create_vertices.restype = ctypes.c_void_p
_create_vertices.argtypes = (
    ctypes.c_char_p, # format
    ctypes.c_uint, # nvertices
    ctypes.c_void_p, # vertices
    ctypes.c_uint, # nindices
    ctypes.c_void_p) # indices

def create_vertices(format, nvertices, vertices, indices):
    return _create_vertices(format, nvertices, _ubyte_addr(vertices), len(indices), _ushort_addr(indices))

create_batch = _engine.rk_create_batch
create_batch.restype = ctypes.c_void_p
create_batch.argtypes = (
    ctypes.c_void_p, # vertices
    ctypes.c_uint, # max_size
    ctypes.c_char_p) # params_format

begin_frame = _engine.rk_begin_frame

select_shader = _engine.rk_select_shader
select_shader.argtypes = (
    ctypes.c_void_p,) # shader

set_input_float = _engine.rk_set_input_float
set_input_float.argtypes = (
    ctypes.c_void_p, # input
    ctypes.c_float) # value

set_input_vec3 = _engine.rk_set_input_vec3
set_input_vec3.argtypes = (
    ctypes.c_void_p, # input
    _vec3_p) # value

set_input_mat3 = _engine.rk_set_input_mat3
set_input_mat3.argtypes = (
    ctypes.c_void_p, # input
    _mat3_p) # value

set_input_mat4 = _engine.rk_set_input_mat4
set_input_mat4.argtypes = (
    ctypes.c_void_p, # input
    _mat4_p) # value

set_param_vec3 = _engine.rk_set_param_vec3
set_param_vec3.argtypes = (
    ctypes.c_uint, # layout
    _vec3_p) # value

select_texture = _engine.rk_select_texture
select_texture.argtypes = (
    ctypes.c_uint, # slot
    ctypes.c_void_p, # texture
    ctypes.c_void_p) # sampler

draw_triangles = _engine.rk_draw_triangles
draw_triangles.argtypes = (
    ctypes.c_void_p,) # triangles

select_vertices = _engine.rk_select_vertices
select_vertices.argtypes = (
    ctypes.c_void_p,) # vertices

draw_batch = _engine.rk_draw_batch
draw_batch.argtypes = (
    ctypes.c_void_p, # batch
    ctypes.c_uint, # size
    ctypes.POINTER(ctypes.c_ubyte), # flags
    ctypes.POINTER(ctypes.c_uint), # meshes
    ctypes.c_void_p) # params

unselect_vertices = _engine.rk_unselect_vertices
unselect_vertices.argtypes = (
    ctypes.c_void_p,) # vertices

unselect_texture = _engine.rk_unselect_texture
unselect_texture.argtypes = (
    ctypes.c_uint, # slot
    ctypes.c_void_p) # texture

unselect_shader = _engine.rk_unselect_shader
unselect_shader.argtypes = (
    ctypes.c_void_p,) # shader

end_frame = _engine.rk_end_frame

destroy_batch = _engine.rk_destroy_batch
destroy_batch.argtypes = (
    ctypes.c_void_p,) # batch

destroy_triangles = _engine.rk_destroy_triangles
destroy_triangles.argtypes = (
    ctypes.c_void_p,) # triangles

destroy_vertices = _engine.rk_destroy_vertices
destroy_vertices.argtypes = (
    ctypes.c_void_p,) # vertices

destroy_texture = _engine.rk_destroy_texture
destroy_texture.argtypes = (
    ctypes.c_void_p,) # texture

destroy_shader = _engine.rk_destroy_shader
destroy_shader.argtypes = (
    ctypes.c_void_p,) # shader

# events.hpp

EVENT_FOCUS_IN = 0
EVENT_FOCUS_OUT = 1
EVENT_KEY_PRESS = 2
EVENT_KEY_RELEASE = 3
EVENT_BUTTON_PRESS = 4
EVENT_BUTTON_RELEASE = 5
EVENT_MOTION = 6

BUTTON_LEFT = 1
BUTTON_MIDDLE = 2
BUTTON_RIGHT = 3
BUTTON_WHEEL_UP = 4
BUTTON_WHEEL_DOWN = 5

class _EventKey(ctypes.Structure):
    _fields_ = (
        ('code', ctypes.c_uint),
        ('symbol', ctypes.c_uint),
        ('character', ctypes.c_wchar))

class _EventButton(ctypes.Structure):
    _fields_ = ('index', ctypes.c_uint),

class _EventMotion(ctypes.Structure):
    _fields_ = (
        ('x', ctypes.c_int),
        ('y', ctypes.c_int))

class _Events(ctypes.Union):
    _fields_ = (
        ('key', _EventKey),
        ('button', _EventButton),
        ('motion', _EventMotion))

class Event(ctypes.Structure):
    _fields_ = ('type', ctypes.c_uint), ('data', _Events)

create_events = _engine.rk_create_events
create_events.restype = ctypes.c_void_p
create_events.argtypes = (
    ctypes.c_void_p,) # display

destroy_events = _engine.rk_destroy_events
destroy_events.argtypes = (
    ctypes.c_void_p, # display
    ctypes.c_void_p) # events

set_key_autorepeat = _engine.rk_set_key_autorepeat
set_key_autorepeat.argtypes = (
    ctypes.c_void_p, # events
    ctypes.c_bool) # autorepeat

set_motion_acceleration = _engine.rk_set_motion_acceleration
set_motion_acceleration.argtypes = (
    ctypes.c_void_p, # events
    ctypes.c_uint, # numerator
    ctypes.c_uint, # denominator
    ctypes.c_uint) # threshold

consume_events = _engine.rk_consume_events
consume_events.restype = ctypes.c_uint
consume_events.argtypes = (
    ctypes.c_void_p, # events
    ctypes.POINTER(Event), # buffer
    ctypes.c_uint) # max_events