package net.halbear.Terrain4J.EngineCore.Logic.Rendering;

import com.beust.jcommander.JCommander;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.ParameterException;
import com.google.gson.FieldNamingPolicy;
import com.google.gson.GsonBuilder;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.MaterialData;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.MeshData;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.ModelBinaryData;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.ModelData;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanUtils;
import org.joml.Vector4f;
import org.lwjgl.PointerBuffer;
import org.lwjgl.assimp.*;
import org.lwjgl.system.MemoryStack;
import org.tinylog.Logger;

import java.io.*;
import java.nio.IntBuffer;
import java.nio.channels.FileChannel;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static org.lwjgl.assimp.Assimp.*;

public class ModelCompiler {
    private static final Pattern EMBED_TEXT_ID = Pattern.compile("\\*([0-9)]+)");
    private static final int FLAGS = aiProcess_GenSmoothNormals | aiProcess_JoinIdenticalVertices
            | aiProcess_Triangulate | aiProcess_FixInfacingNormals | aiProcess_CalcTangentSpace
            | aiProcess_PreTransformVertices;

    @Parameter(names = "-model", description = "Model Path", required = true)
    private String ModelPath;

    public static void main(String[] args){
        var Main = new ModelCompiler();
        var jCommander = JCommander.newBuilder().addObject(Main).build();
        try{
            jCommander.parse(args);
            Main.Process();
        } catch(ParameterException exception){
            jCommander.usage();
        } catch(IOException exception){
            Logger.error("Error Compiling model",exception);
        }
    }
    public static void Process(String ModelPath) throws IOException{
        Logger.debug("Loading 3D Model Data from [{}]", ModelPath);
        var ModelSourceFile = new File(ModelPath);
        if (!ModelSourceFile.exists()) throw new RuntimeException("Model path does not exist->["+ModelPath+"]");
        AIScene aiScene = aiImportFile(ModelPath, FLAGS);
        if (aiScene == null) throw new RuntimeException("Model ["+ModelPath+"] could not be loaded");

        String ModelID = ModelSourceFile.getName();
        if(ModelID.contains(".")) ModelID = ModelID.substring(0, ModelID.lastIndexOf('.'));

        ModelBinaryData BinaryData = new ModelBinaryData(ModelPath);

        int MaterialCount = aiScene.mNumMaterials();
        Logger.debug("Number of materials: {}",MaterialCount);
        List<MaterialData> MaterialList = new ArrayList<>();
        File ParentDir = ModelSourceFile.getParentFile();
        for(int i = 0; i < MaterialCount; i++){
            var aiMaterial = AIMaterial.create(aiScene.mMaterials().get(i));
            MaterialData material = ProcessMaterial(aiScene, aiMaterial, ModelID, ParentDir.getPath(), i);
            MaterialList.add(material);
        }

        int MeshCount = aiScene.mNumMeshes();
        PointerBuffer aiMeshes = aiScene.mMeshes();
        List<MeshData> meshList = new ArrayList<>();
        for(int i = 0; i < MeshCount; i++){
            AIMesh aiMesh = AIMesh.create(aiMeshes.get(i));
            MeshData newMesh = ProcessMesh(aiMesh, MaterialList, i,BinaryData);
            meshList.add(newMesh);
        }

        var model = new ModelData(ModelID, meshList, BinaryData.GetVertexFilePath(), BinaryData.GetIndexFilePath());

        String CreateModelFile = ModelPath.substring(0,ModelPath.lastIndexOf('.')) + ".json";
        Writer writer = new FileWriter(CreateModelFile);
        var gson = new GsonBuilder().setFieldNamingPolicy(FieldNamingPolicy.UPPER_CAMEL_CASE).setPrettyPrinting().create();
        gson.toJson(model, writer);
        writer.flush();
        writer.close();
        Logger.info("Generated Model File [{}]", CreateModelFile);

        String CreateMaterialFile = ModelPath.substring(0,ModelPath.lastIndexOf('.')) + "_mat.json";
        writer = new FileWriter(CreateMaterialFile);
        gson.toJson(MaterialList,writer);
        writer.flush();
        writer.close();
        BinaryData.Close();
        aiReleaseImport(aiScene);
        Logger.info("Generated Vertex File [{}] and Index File [{}]",BinaryData.GetVertexFilePath(),BinaryData.GetIndexFilePath());
    }

    private void Process() throws IOException{
        Process(ModelPath);
    }

    public static MaterialData ProcessMaterial(AIScene aiScene, AIMaterial aiMaterial, String ModelName, String BaseDirectory, int Position)  throws IOException{
        Vector4f diffuse = new Vector4f();
        AIColor4D colour = AIColor4D.create();

        int Result = aiGetMaterialColor(aiMaterial, AI_MATKEY_COLOR_DIFFUSE, aiTextureType_NONE, 0, colour);
        if(Result == aiReturn_SUCCESS){
            diffuse.set(colour.r(), colour.g(), colour.b(), colour.a());
        }
        String DiffuseTexture = ProcessTexture(aiScene, aiMaterial, BaseDirectory, aiTextureType_DIFFUSE);
        return new MaterialData(ModelName + "-mat-" + Position,DiffuseTexture,diffuse);
    }

    public static MeshData ProcessMesh(AIMesh aiMesh, List<MaterialData> MaterialList, int MeshPosition, ModelBinaryData BinaryData) throws IOException {
        List<Float> vertices = ProcessVertices(aiMesh);
        List<Float> textureCoordinates = ProcessTextureCoords(aiMesh);
        List<Integer> indices = ProcessIndices(aiMesh);

        int VerticesSize =  vertices.size();
        if(textureCoordinates.isEmpty()){textureCoordinates = Collections.nCopies((VerticesSize/3)*2,0.0f);}

        DataOutputStream VerticesOutput = BinaryData.GetVertexOutput();
        int Rows = VerticesSize/3;
        int VertexIncrement = (VerticesSize + textureCoordinates.size()) * VulkanUtils.FLOAT_SIZE;
        for(int row = 0; row < Rows; row++){
            int StartPos = row * 3;
            int StartTextureCoord = row * 2;
            VerticesOutput.writeFloat(vertices.get(StartPos));
            VerticesOutput.writeFloat(vertices.get(StartPos + 1));
            VerticesOutput.writeFloat(vertices.get(StartPos + 2));
            VerticesOutput.writeFloat(textureCoordinates.get(StartTextureCoord));
            VerticesOutput.writeFloat(textureCoordinates.get(StartTextureCoord + 1));
        }

        DataOutputStream IndexOutput = BinaryData.GetIndexOutput();
        int IndexSize = indices.size();
        int IndexIncrement = IndexSize * VulkanUtils.INT_SIZE;
        for(int Index = 0; Index < IndexSize; Index++){
            IndexOutput.writeInt(indices.get(Index));
        }

        String MeshID = aiMesh.mName().dataString() + "_" + MeshPosition;
        int MaterialIndex = aiMesh.mMaterialIndex();
        String MaterialID = "";
        if(MaterialIndex >= 0 && MaterialIndex < MaterialList.size()){
            MaterialID = MaterialList.get(MaterialIndex).ID();
        }

        var meshData = new MeshData(MeshID, MaterialID, BinaryData.GetVertexOffset(), VertexIncrement,BinaryData.GetIndexOffset(), IndexIncrement);
        BinaryData.IncrementVertexOffset(VertexIncrement);
        BinaryData.IncrementIndexOffset(IndexIncrement);
        return meshData;
    }

    public  static List<Float> ProcessTextureCoords(AIMesh aiMesh){
        List<Float> TextureCoords = new ArrayList<>();
        AIVector3D.Buffer aiTextureCoords = aiMesh.mTextureCoords(0);
        int CoordinateCount = aiTextureCoords != null ? aiTextureCoords.remaining() : 0;
        for(int i = 0; i < CoordinateCount; i++){
            AIVector3D TextureCoord = aiTextureCoords.get();
            TextureCoords.add(TextureCoord.x());
            TextureCoords.add(1 - TextureCoord.y());
        }
        return TextureCoords;
    }

    public static List<Integer> ProcessIndices(AIMesh aiMesh){
        List<Integer> Indices = new ArrayList<>();
        int FaceCount = aiMesh.mNumFaces();
        AIFace.Buffer aiFaces = aiMesh.mFaces();
        for(int i = 0; i < FaceCount; i++){
            AIFace aiFace = aiFaces.get(i);
            IntBuffer buffer = aiFace.mIndices();
            while(buffer.remaining() > 0){
                Indices.add(buffer.get());
            }
        }
        return Indices;
    }

    public static List<Float> ProcessVertices(AIMesh aiMesh){
        List<Float> vertices = new ArrayList<>();
        AIVector3D.Buffer aiVertices = aiMesh.mVertices();
        while(aiVertices.remaining() > 0){
            AIVector3D aiVertex = aiVertices.get();
            vertices.add(aiVertex.x());
            vertices.add(aiVertex.y());
            vertices.add(aiVertex.z());
        }
        return vertices;
    }

    public static String ProcessTexture(AIScene aiScene, AIMaterial aiMaterial, String BaseDirectory,int TextType ) throws IOException{
        String TexturePath;
        try(var MemStack = MemoryStack.stackPush()){
            int EmbeddedTextureCount = aiScene.mNumTextures();
            AIString aiTexturePath = AIString.calloc(MemStack);
            aiGetMaterialTexture(aiMaterial,TextType,0,aiTexturePath,(IntBuffer) null,null,null,null, null,null);
            TexturePath = aiTexturePath.dataString();
            if(TexturePath != null && !TexturePath.isEmpty()){
                Matcher matcher = EMBED_TEXT_ID.matcher(TexturePath);
                int EmbeddedTextureIndex = matcher.matches() && matcher.groupCount() > 0 ? Integer.parseInt(matcher.group(1)) : -1;
                if (EmbeddedTextureIndex >= 0 && EmbeddedTextureIndex < EmbeddedTextureCount){
                    var aiTexture = AITexture.create(aiScene.mTextures().get(EmbeddedTextureIndex));
                    String BaseFile = aiTexture.mFilename().dataString() + ".png";
                    TexturePath = BaseDirectory + File.separator + BaseFile;
                    Logger.info("Dumping Texture File to [{}]",TexturePath);
                    var Channel = FileChannel.open(Path.of(TexturePath), StandardOpenOption.CREATE,StandardOpenOption.WRITE);
                    Channel.write(aiTexture.pcDataCompressed());
                } else{
                    TexturePath = BaseDirectory + File.separator + new File(TexturePath).getName();
                }
            }
            return TexturePath;
        }
    }
}
