# 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")

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

    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 set(self, x, y, z, w):
        self.x = x
        self.y = y
        self.z = z
        self.w = w

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

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

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)

def _flag(x):
    return 1 << x

TEXTURE_FORMAT_SRGB8_A8 = 0
TEXTURE_FORMAT_RGBA8 = 1
TEXTURE_FORMAT_RGB10_A2 = 2
TEXTURE_FORMAT_32F = 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_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

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

_vec3p = ctypes.POINTER(vec3)
_vec4p = ctypes.POINTER(vec4)
_mat3p = ctypes.POINTER(mat3)
_mat4p = ctypes.POINTER(mat4)

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

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

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

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

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


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

mat3_rotation = _engine.rk_mat3_rotation
mat3_rotation.argtypes = (
    _mat3p, # ret
    _vec3p, # axis
    ctypes.c_float) # angle

mat3_mul_vec3 = _engine.rk_mat3_mul_vec3
mat3_mul_vec3.argtypes = (
    _vec3p, # ret
    _mat3p, # a
    _vec3p) # b

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

mat4_projection = _engine.rk_mat4_projection
mat4_projection.argtypes = (
    _mat4p, # 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 = (
    _mat4p, # ret
    _vec3p, # position
    _vec3p) # lookat

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

mat4_mul_vec4 = _engine.rk_mat4_mul_vec4
mat4_mul_vec4.argtypes = (
    _vec4p, # ret
    _mat4p, # a
    _vec4p) # b

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

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

_load_shader = _engine.rk_load_shader
_load_shader.restype = ctypes.c_void_p
_load_shader.argtypes = (
    ctypes.c_uint, # vert_nlines
    ctypes.c_void_p, # vert_lines
    ctypes.c_uint, # frag_nlines
    ctypes.c_void_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, ctypes.addressof(vert_lines),
        frag_nlines, ctypes.addressof(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, _voidp(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, _floatp(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, _ubytep(vertices), len(indices), _ushortp(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
    _vec3p) # value

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

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

set_param_vec3 = _engine.rk_set_param_vec3
set_param_vec3.argtypes = (
    ctypes.c_uint, # layout
    _vec3p) # 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

terminate = _engine.rk_terminate

EVENT_KEY_PRESS = 0
EVENT_KEY_RELEASE = 1

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

class _Events(ctypes.Union):
    _fields_ = (
        ('key', _EventKey),
    )

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

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

flush_events = _engine.rk_flush_events