# 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 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,
    engine.VERTEX_FORMAT_VEC3_INT10,
    engine.VERTEX_FORMAT_VEC2_USHORT)

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, 'B')
    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, 'B', (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

    @classmethod
    def from_archive(cls, file):
        _read_magic(file, b'TX')
        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'TX')
        _write_string(file, self.name)
        _write_struct(file, 'IIIII', (self.format, self.width, self.height, self.nlevels, self.flags))
        _write_blob(file, self.pixels)

def create_texture(slot, input, data):
    return engine.create_texture(
        slot, input, data.format, data.width, data.height, data.nlevels, data.flags, data.pixels)

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

    def __init__(self, name, vertices, indices):
        if len(vertices) % VERTEX_SIZE != 0:
            raise RuntimeError("Vertex format mismatch!")
        self.name = name
        self.vertices = vertices
        self.indices = indices

    @classmethod
    def from_archive(cls, file):
        _read_magic(file, b'VT')
        name = _read_string(file)
        nvertices, nindices = _read_struct(file, 'II')
        vertices = _read_array(file, 'B', nvertices * VERTEX_SIZE)
        indices = _read_array(file, 'H', nindices)
        return cls(name, vertices, indices)

    def to_archive(self, file):
        _write_magic(file, b'VT')
        _write_string(file, self.name)
        _write_struct(file, 'II', (len(self.vertices) // VERTEX_SIZE, len(self.indices)))
        _write_array(file, self.vertices)
        _write_array(file, self.indices)

def create_vertices(data):
    return engine.create_vertices(VERTEX_FORMAT, len(data.vertices) // VERTEX_SIZE, data.vertices, data.indices)

class ModelData:
    __slots__ = 'name', 'flags', 'texlevel', 'mesh'

    def __init__(self, name, flags, texlevel, mesh):
        self.name = name
        self.flags = flags
        self.texlevel = texlevel
        self.mesh = mesh

    @classmethod
    def from_archive(cls, file):
        _read_magic(file, b'MD')
        name = _read_string(file)
        flags, texlevel, mesh = _read_struct(file, 'BHI')
        return ModelData(name, flags, texlevel, mesh)

    def to_archive(self, file):
        _write_magic(file, b'MD')
        _write_string(file, self.name)
        _write_struct(file, 'BHI', (self.flags, self.texlevel, self.mesh))

class Model:
    __slots__ = 'flags', 'texlevel', 'mesh'

    def __init__(self, flags, texlevel, mesh):
        self.flags = flags
        self.texlevel = texlevel
        self.mesh = mesh

    def spawn(self, batch, translation = engine.vec3_zero, orientation = engine.vec3_forward):
        return batch.append(self.flags, self.texlevel, self.mesh, translation, orientation)

def create_model(data):
    return Model(data.flags, data.texlevel, data.mesh)

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

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

    def destroy(self):
        for vertices in self.vertices_db.values():
            engine.destroy_vertices(vertices)
        for texture in self.textures_db.values():
            engine.destroy_texture(texture)
        self.textures_db.clear()
        self.vertices_db.clear()
        self.models_db.clear()

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

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

    def get_model(self, name):
        return self.models_db[name]

    @classmethod
    def _new_texture(cls, data):
        return data

    @classmethod
    def _new_vertices(cls, data):
        return data

    @classmethod
    def _new_model(cls, data):
        return data

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

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

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

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

class RuntimeArchive(Archive):
    @classmethod
    def _new_texture(cls, data):
        return create_texture(0, b'u_texture_sampler', data)

    @classmethod
    def _new_vertices(cls, data):
        return create_vertices(data)

    @classmethod
    def _new_model(cls, data):
        return create_model(data)