# 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
import struct
from array import array
from pathlib import Path

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

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_VEC2_FLOAT = 1
VERTEX_FORMAT_VEC2_USHORT = 2
VERTEX_FORMAT_VEC3_FLOAT = 3
VERTEX_FORMAT_VEC3_INT10 = 4

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

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

BATCH_MAX_SIZE = 65536

BATCH_TRANSLATION_FORMAT_FLOAT = 0
BATCH_TRANSLATION_FORMAT_SHORT = 1

BATCH_ORIENTATION_FORMAT_NONE = 0
BATCH_ORIENTATION_FORMAT_FLOAT = 1
BATCH_ORIENTATION_FORMAT_INT10 = 2

#TODO: remove from engine
vec2_zero = (0.0, 0.0)
vec3_zero = (0.0, 0.0, 0.0)
vec3_right = (1.0, 0.0, 0.0)
vec3_forward = (0.0, 1.0, 0.0)
vec3_up = (0.0, 0.0, 1.0)
vec4_zero = (0.0, 0.0, 0.0, 1.0)
mat3_identity = (
    1.0, 0.0, 0.0,
    0.0, 1.0, 0.0,
    0.0, 0.0, 1.0)
mat4_identity = (
    1.0, 0.0, 0.0, 0.0,
    0.0, 1.0, 0.0, 1.0,
    0.0, 0.0, 1.0, 0.0,
    0.0, 0.0, 0.0, 1.0)

def vec2(v = vec2_zero):
    assert len(v) == 2
    return array('f', v)

def vec3(v = vec3_zero):
    assert len(v) == 3
    return array('f', v)

def vec4(v = vec4_zero):
    assert len(v) == 4
    return array('f', v)

def mat3(m = mat3_identity):
    assert len(m) == 9
    return array('f', m)

def mat4(m = mat4_identity):
    assert len(m) == 16
    return array('f', m)

_vec2_t = (ctypes.c_float * 2)
_vec2 = _vec2_t.from_buffer

_vec3_t = (ctypes.c_float * 3)
_vec3 = _vec3_t.from_buffer

_vec4_t = (ctypes.c_float * 4)
_vec4 = _vec4_t.from_buffer

_mat3_t = (ctypes.c_float * 9)
_mat3 = _mat3_t.from_buffer

_mat4_t = (ctypes.c_float * 16)
_mat4 = _mat4_t.from_buffer

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 _uintp(x):
    assert x.typecode == 'I'
    return x.buffer_info()[0]

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

_mat3_rotation = _lib.rk_mat3_rotation
_mat3_rotation.argtypes = (
    _mat3_t, # ret
    _vec3_t, # axis
    ctypes.c_float) # angle

def mat3_rotation(ret, axis, angle):
    assert len(ret) == 9 and len(axis) == 3
    _mat3_rotation(_mat3(ret), _vec3(axis), angle)

_mat3_mul_vec3 = _lib.rk_mat3_mul_vec3
_mat3_mul_vec3.argtypes = (
    _vec3_t, # ret
    _mat3_t, # a
    _vec3_t) # b

def mat3_mul_vec3(ret, a, b):
    assert len(ret) == 3 and len(a) == 9 and len(b) == 3
    _mat3_mul_vec3(_vec3(ret), _mat3(a), _vec3(b))

_mat3_mul_mat3 = _lib.rk_mat3_mul_mat3
_mat3_mul_mat3.argtypes = (
    _mat3_t, # ret
    _mat3_t, # a
    _mat3_t) # b

def mat3_mul_mat3(ret, a, b):
    assert len(ret) == 9 and len(a) == 9 and len(b) == 9
    _mat3_mul_mat3(_mat3(ret), _mat3(a), _mat3(b))

_mat4_projection = _lib.rk_mat4_projection
_mat4_projection.argtypes = (
    _mat4_t, # ret
    ctypes.c_float, # hfov
    ctypes.c_float, # ratio
    ctypes.c_float, # near
    ctypes.c_float) # far

def mat4_projection(ret, hfov, ratio, near, far):
    assert len(ret) == 16
    _mat4_projection(_mat4(ret), hfov, ratio, near, far)

_mat4_lookat = _lib.rk_mat4_lookat
_mat4_lookat.argtypes = (
    _mat4_t, # ret
    _vec3_t, # position
    _vec3_t) # lookat

def mat4_lookat(ret, position, lookat):
    assert len(ret) == 16 and len(position) == 3 and len(lookat) == 3
    _mat4_lookat(_mat4(ret), _vec3(position), _vec3(lookat))

_mat4_mul_vec3 = _lib.rk_mat4_mul_vec3
_mat4_mul_vec3.argtypes = (
    _vec3_t, # ret
    _mat4_t, # a
    _vec3_t, # b
    ctypes.c_float) # w

def mat4_mul_vec3(ret, a, b, w):
    assert len(ret) == 3 and len(a) == 16 and len(b) == 3
    _mat4_mul_vec3(_vec3(ret), _mat4(a), _vec3(b), w)

_mat4_mul_vec4 = _lib.rk_mat4_mul_vec4
_mat4_mul_vec4.argtypes = (
    _vec4_t, # ret
    _mat4_t, # a
    _vec4_t) # b

def mat4_mul_vec4(ret, a, b):
    assert len(ret) == 4 and len(a) == 16 and len(b) == 4
    _mat4_mul_vec4(_vec4(ret), _mat4(a), _vec4(b))

_mat4_mul_mat4 = _lib.rk_mat4_mul_mat4
_mat4_mul_mat4.argtypes = (
    _mat4_t, # ret
    _mat4_t, # a
    _mat4_t) # b

def mat4_mul_mat4(ret, a, b):
    assert len(ret) == 16 and len(a) == 16 and len(b) == 16
    _mat4_mul_mat4(_mat4(ret), _mat4(a), _mat4(b))

initialize = _lib.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 = _lib.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))

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

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

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

_set_input_vec3 = _lib.rk_set_input_vec3
_set_input_vec3.argtypes = (
    ctypes.c_void_p, # input
    _vec3_t) # value

def set_input_vec3(input, value):
    assert len(value) == 3
    _set_input_vec3(input, _vec3(value))

_set_input_mat3 = _lib.rk_set_input_mat3
_set_input_mat3.argtypes = (
    ctypes.c_void_p, # input
    _mat3_t) # value

def set_input_mat3(input, value):
    assert len(value) == 9
    _set_input_mat3(input, _mat3(value))

_set_input_mat4 = _lib.rk_set_input_mat4
_set_input_mat4.argtypes = (
    ctypes.c_void_p, # input
    _mat4_t) # value

def set_input_mat4(input, value):
    assert len(value) == 16
    _set_input_mat4(input, _mat4(value))

_create_texture = _lib.rk_create_texture
_create_texture.restype = ctypes.c_void_p
_create_texture.argtypes = (
    ctypes.c_uint, # slot
    ctypes.c_char_p, # input
    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(slot, input, 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(slot, input, format, width, height, nlevels, flags, _voidp(pixels))

_create_triangles = _lib.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 = _lib.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 = _lib.rk_create_batch
create_batch.restype = ctypes.c_void_p
create_batch.argtypes = (
    ctypes.c_uint, # max_size
    ctypes.c_uint, # translation_format
    ctypes.c_uint) # orientation_format

begin_frame = _lib.rk_begin_frame

select_texture = _lib.rk_select_texture
select_texture.argtypes = (
    ctypes.c_void_p,) # texture

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

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

_draw_batch = _lib.rk_draw_batch
_draw_batch.argtypes = (
    ctypes.c_void_p, # batch
    ctypes.c_uint, # size
    ctypes.c_void_p, # flags
    ctypes.c_void_p, # texlevels
    ctypes.c_void_p, # meshes
    ctypes.c_void_p, # translations
    ctypes.c_void_p) # orientations

def draw_batch(batch, flags, texlevels, meshes, translations, orientations):
    size = len(flags)
    assert len(texlevels) == size and len(meshes) == size and len(translations) == size * 3
    assert not orientations or len(orientations) == size * 3
    _draw_batch(batch, size, _ubytep(flags), _ushortp(texlevels), _uintp(meshes), _floatp(translations),
        _floatp(orientations) if orientations else None)

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

unselect_texture = _lib.rk_unselect_texture
unselect_texture.argtypes = (
    ctypes.c_void_p,) # texture

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

end_frame = _lib.rk_end_frame

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

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

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

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

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

terminate = _lib.rk_terminate