# 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 sys
import struct
from array import array
from pathlib import Path
from itertools import starmap, chain, repeat

from engine import *

from game.math import float_s8, vec3_srgb8a8
from game.resources import load_png, TextureData, VerticesData, ModelData, Archive

_texture_flags = TEXTURE_FLAG_MIPMAPS | TEXTURE_FLAG_MIN_LINEAR | TEXTURE_FLAG_MAG_NEAREST

class ObjArchive(Archive):
    def __init__(self):
        super().__init__()

    def import_mtl(self, mtlpath):
        def load_value(x):
            x = round(float(x), 6)
            if x == -0.0:
                x = 0.0
            return x
        size = None
        def load_texture(pngpath, nchannels):
            width, height, pixels = load_png(pngpath)
            assert pixels.typecode == 'B'
            assert (size is None or size == (width, height))
            assert len(pixels) == width * height * nchannels
            return ((width, height), pixels)
        print("Importing mtl", mtlpath)
        texture = array('B')
        texlevels = {}
        name = ''
        texlevel = 0
        color = (0, 0, 0, 0)
        metallic = 0
        specular = 0
        roughness = 0
        bump = 0
        def finish(_texture, _color, _metallic, _specular, _roughness, _bump):
            width, height = size
            if not isinstance(_color, array):
                _color = array('B', _color * (width * height))
            _texture.extend(_color)
            if not isinstance(_metallic, array):
                _metallic = repeat(_metallic, width * height)
            if not isinstance(_specular, array):
                _specular = repeat(_specular, width * height)
            if not isinstance(_roughness, array):
                _roughness = repeat(_roughness, width * height)
            if not isinstance(_bump, array):
                _bump = repeat(_bump, width * height)
            _texture.extend(chain.from_iterable(zip(_metallic, _specular, _roughness, _bump)))
        for line in open(mtlpath):
            data = line.split()
            if not data:
                continue
            code = data.pop(0)
            if code == '#':
                continue
            elif code == 'newmtl':
                if name:
                    if size is None:
                        size = (1, 1)
                    finish(texture, color, metallic, specular, roughness, bump)
                    texlevels[name] = texlevel
                    texlevel += 2
                name = data[0]
                print("Importing material", name)
                color = (0, 0, 0, 0)
                metallic = 0
                specular = 0
                roughness = 0
                bump = 0
            elif code == 'Kd': # color
                color = vec3_srgb8a8(tuple(map(load_value, data)))
                assert len(color) == 4
            elif code == 'Ni': # ior -> metallic
                metallic = float_s8(load_value(data[0]))
            elif code == 'Ks': # specular
                specular = float_s8(load_value(data[0]))
            elif code == 'Ns': # roughness
                roughness = float_s8(1.0 - load_value(data[0]) / 1000.0)
            elif code == 'map_Kd':
                pngpath = mtlpath.parent / data[0]
                print("Importing color texture", pngpath)
                size, color = load_texture(pngpath, 4)
                pngpath = pngpath.parent / (pngpath.stem + '_b.png')
                if pngpath.exists():
                    print("Importing bump channel", pngpath)
                    size, bump = load_texture(pngpath, 1)
            elif code == 'map_Ni':
                pngpath = mtlpath.parent / data[0]
                print("Importing metallic channel", pngpath)
                size, metallic = load_texture(pngpath, 1)
            elif code == 'map_Ks':
                pngpath = mtlpath.parent / data[0]
                print("Importing specular channel", pngpath)
                size, specular = load_texture(pngpath, 1)
            elif code == 'map_Ns':
                pngpath = mtlpath.parent / data[0]
                print("Importing roughness channel", pngpath)
                size, roughness = load_texture(pngpath, 1)
        if name:
            if size is None:
                size = (1, 1)
            finish(texture, color, metallic, specular, roughness, bump)
            texlevels[name] = texlevel
            texlevel += 2
        name = str(mtlpath.stem)
        assert texlevel < 255
        assert name not in self.textures_db.keys()
        print("Storing texture", name)
        width, height = size
        self.textures_db[name] = TextureData(
            name, TEXTURE_FORMAT_SRGB8_A8, width, height, texlevel, _texture_flags, texture)
        return texlevels

    def import_obj(self, objpath):
        def load_coord(x):
            x = round(float(x), 6)
            if x == -0.0:
                x = 0.0
            return x
        def load_index(x):
            return int(x) - 1
        print("Importing obj", objpath)
        texlevels = {}
        name = ''
        texlevel = 0
        mesh = []
        positions = []
        normals = []
        texcoords = []
        objects = {}
        for line in open(objpath):
            data = line.split()
            if not data:
                continue
            code = data.pop(0)
            if code == '#':
                continue
            elif code == 'mtllib':
                texlevels = self.import_mtl(objpath.parent / data[0])
            elif code == 'o':
                if name:
                    assert mesh
                    objects[name] = mesh
                name = data[0]
                print("Importing object", name)
                mesh = []
            elif code == 'v':
                position = tuple(map(load_coord, data))
                assert len(position) == 3
                positions.append(position)
            elif code == 'vn':
                normal = tuple(map(load_coord, data))
                assert len(normal) == 3
                normals.append(normal)
            elif code == 'vt':
                texcoord = tuple(map(load_coord, data))
                assert len(texcoord) == 2
                texcoords.append((texcoord[0], 1.0 - texcoord[1]))
            elif code == 'usemtl':
                texlevel = texlevels[data[0]]
            elif code == 'f':
                indices = tuple(map(lambda x: tuple(map(load_index, x.split('/'))), data))
                assert len(indices) == 3
                assert len(indices[0]) == 3
                assert len(indices[1]) == 3
                assert len(indices[2]) == 3
                triangle = tuple(map(
                    lambda x: positions[x[0]] + normals[x[2]] + texcoords[x[1]] + (texlevel,), indices))
                assert len(triangle) == 3 and len(set(triangle)) == 3
                mesh.append(triangle)
        if name:
            assert mesh
            objects[name] = mesh
        vertices = set()
        for mesh in objects.values():
            vertices |= frozenset(chain.from_iterable(mesh))
        vertices = tuple(vertices)
        indices = []
        models = {}
        for name, mesh in sorted(objects.items()):
            offset = len(indices)
            assert offset < 65536
            indices.extend(map(vertices.index, chain.from_iterable(mesh)))
            count = len(indices) - offset
            assert count % 3 == 0
            count //= 3
            assert count < 65536
            print(name, ": offset =", offset, "triangles =", count)
            models[name] = (offset, count)
        name = str(objpath.stem)
        assert name not in self.vertices_db.keys()
        #TODO: move to math
        def pack_10(_x):
            return round(_x * (512.0 if _x < 0.0 else 511.0)) & 1023
        def pack_u10(_x):
            return round(_x * 1023.0) & 1023
        def pack_vertex(_px, _py, _pz, _nx, _ny, _nz, _s, _t, _tl):
            n = (pack_10(_nz) << 20) | (pack_10(_ny) << 10) | pack_10(_nx)
            t = ((_tl & 1023) << 20) | (pack_u10(_t) << 10) | pack_u10(_s)
            return struct.pack('fffII', _px, _py, _pz, n, t)
        self.vertices_db[name] = VerticesData(name,
            array('B', b''.join(starmap(pack_vertex, vertices))),
            array('H', indices))
        for name, (offset, count) in sorted(models.items()):
            if name[0] != '_':
                print("Storing", name)
                assert name not in self.models_db.keys()
                self.models_db[name] = ModelData(
                    name, INSTANCE_FLAG_VISIBLE, offset | count << 16)

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("Usage: python3 -m tools.obj2rkar output.rkar input.obj ...")
        sys.exit(2)
    importer = ObjArchive()
    for argv in sys.argv[2:]:
        objpath = Path(argv)
        importer.import_obj(objpath)
    outpath = Path(sys.argv[1])
    print("Exporting", outpath)
    importer.save(outpath)
    print("Done.")
