# 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/>.

from itertools import chain
import struct
from array import array
from pathlib import Path
import png
import engine

VERTEX_SIZE = 20
VERTEX_FORMAT = engine.vertex_format(
    engine.VERTEX_FORMAT_VEC3_FLOAT, # position
    engine.VERTEX_FORMAT_VEC3_INT10 | engine.VERTEX_FORMAT_NORMALIZE, # normal
    engine.VERTEX_FORMAT_VEC3_UINT10) # texcoords

def load_png(path):
    width, height, data, _ = png.Reader(filename = path).read_flat()
    return (width, height, data)

def _read_magic(file, magic):
    assert magic
    if file.read(len(magic)) != magic:
        raise RuntimeError("Archive magic mismatch!", magic)

def _write_magic(file, magic):
    assert magic
    size = file.write(magic)
    assert size == len(magic)

def _read_struct(file, format):
    assert format
    size = struct.calcsize(format)
    assert size
    data = file.read(size)
    assert len(data) == size
    _read_magic(file, b'RK')
    return struct.unpack(format, data)

def _write_struct(file, format, *elems):
    assert format
    data = struct.pack(format, *elems)
    size = file.write(data)
    assert size == len(data)
    _write_magic(file, b'RK')

def _read_string(file):
    length, = _read_struct(file, 'I')
    assert length
    data = file.read(length)
    assert len(data) == length
    _read_magic(file, b'RK')
    return str(data, encoding='ascii')

def _write_string(file, string):
    data = bytes(string, encoding='ascii')
    assert data and len(data) < 256
    _write_struct(file, 'I', len(data))
    size = file.write(data)
    assert size == len(data)
    _write_magic(file, b'RK')

def _read_array(file, format, length):
    assert format
    data = array(format)
    data.fromfile(file, length)
    assert len(data) == length
    _read_magic(file, b'RK')
    return data

def _write_array(file, array):
    assert array
    array.tofile(file)
    _write_magic(file, b'RK')

def _read_blob(file):
    typecode, length = _read_struct(file, 'II')
    typecode = chr(typecode)
    assert typecode, length
    data = array(typecode)
    data.fromfile(file, length)
    assert len(data) == length
    _read_magic(file, b'RK')
    return data

def _write_blob(file, array):
    assert array
    _write_struct(file, 'II', ord(array.typecode), len(array))
    array.tofile(file)
    _write_magic(file, b'RK')

class TextureData:
    __slots__ = 'name', 'format', 'width', 'height', 'nlevels', 'flags', 'pixels'

    def __init__(self, name, format, width, height, nlevels, flags, pixels):
        assert pixels.typecode == engine.TEXTURE_FORMAT_TYPECODE[format]
        assert len(pixels) == width * height * max(1, nlevels) * engine.TEXTURE_FORMAT_NELEMS[format]
        self.name = name
        self.format = format
        self.width = width
        self.height = height
        self.nlevels = nlevels
        self.flags = flags
        self.pixels = pixels

    def __iter__(self):
        yield self.format
        yield self.width
        yield self.height
        yield self.nlevels
        yield self.flags
        yield self.pixels

    @classmethod
    def from_archive(cls, file):
        _read_magic(file, b'TXTR')
        name = _read_string(file)
        format, width, height, nlevels, flags = _read_struct(file, 'IIIII')
        pixels = _read_blob(file)
        assert pixels.typecode == engine.TEXTURE_FORMAT_TYPECODE[format]
        assert len(pixels) == width * height * max(1, nlevels) * engine.TEXTURE_FORMAT_NELEMS[format]
        return cls(name, format, width, height, nlevels, flags, pixels)

    def to_archive(self, file):
        _write_magic(file, b'TXTR')
        _write_string(file, self.name)
        _write_struct(file, 'IIIII', self.format, self.width, self.height, self.nlevels, self.flags)
        _write_blob(file, self.pixels)

class VerticesData:
    __slots__ = 'name', 'vertices', 'indices', 'meshes', 'meshes_db'

    def __init__(self, name, vertices, indices, meshes):
        if len(vertices) % VERTEX_SIZE != 0:
            raise RuntimeError("Vertex format mismatch!")
        self.name = name
        self.vertices = vertices
        self.indices = indices
        self.meshes = meshes.values()
        self.meshes_db = dict(zip(meshes.keys(), range(len(meshes))))

    def __iter__(self):
        yield VERTEX_FORMAT
        yield len(self.vertices) // VERTEX_SIZE
        yield self.vertices
        yield self.indices
        yield array('I', chain.from_iterable(self.meshes))

    def get_mesh(self, name):
        return self.meshes_db[name]

    @classmethod
    def from_archive(cls, file):
        _read_magic(file, b'VERT')
        name = _read_string(file)
        nvertices, nindices, nmeshes = _read_struct(file, 'III')
        vertices = _read_array(file, 'B', nvertices * VERTEX_SIZE)
        indices = _read_array(file, 'H', nindices)
        names = [_read_string(file) for _ in range(nmeshes)]
        meshes = [_read_struct(file, 'II') for _ in range(nmeshes)]
        return cls(name, vertices, indices, dict(zip(names, meshes)))

    def to_archive(self, file):
        _write_magic(file, b'VERT')
        _write_string(file, self.name)
        _write_struct(file, 'III', len(self.vertices) // VERTEX_SIZE, len(self.indices), len(self.meshes_db))
        _write_array(file, self.vertices)
        _write_array(file, self.indices)
        for name in self.meshes_db.keys():
            _write_string(file, name)
        for mesh in self.meshes:
            _write_struct(file, 'II', *mesh)

class Archive:
    __slots__ = 'textures_db', 'vertices_db'

    def __init__(self, textures_db = None, vertices_db = None):
        self.textures_db = textures_db or {}
        self.vertices_db = vertices_db or {}

    def get_texture(self, name):
        return self.textures_db[name]

    def get_vertices(self, name):
        return self.vertices_db[name]

    @classmethod
    def from_archive(cls, file):
        textures_db = {}
        vertices_db = {}
        _read_magic(file, b'RKAR')
        ntextures, nvertices = _read_struct(file, 'II')
        for _ in range(ntextures):
            data = TextureData.from_archive(file)
            textures_db[data.name] = data
        for _ in range(nvertices):
            data = VerticesData.from_archive(file)
            vertices_db[data.name] = data
        return cls(textures_db, vertices_db)

    def to_archive(self, file):
        _write_magic(file, b'RKAR')
        _write_struct(file, 'II', len(self.textures_db), len(self.vertices_db))
        for _, data in self.textures_db.items():
            data.to_archive(file)
        for _, data in self.vertices_db.items():
            data.to_archive(file)

    @classmethod
    def load(cls, filename):
        file = open(Path(filename), 'rb')
        archive = cls.from_archive(file)
        file.close()
        return archive

    def save(self, filename):
        file = open(Path(filename), 'wb')
        self.to_archive(file)
        file.close()