#include <EGL/egl.h>
#include <EGL/eglext.h>
#include <GLES3/gl3.h>
#include <png.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <limits>
#include <string>
#include <string_view>
#include <vector>

namespace
{
    struct FVertex
    {
        float Position[3];
        float Normal[3];
        float Height;
    };

    struct FScene
    {
        std::uint32_t VoxelSizeMicrons=0;
        std::array<double,3> CameraMeters{};
        double YawDegrees=0.0,PitchDegrees=0.0;
        std::vector<FVertex> Vertices;
        std::vector<std::uint32_t> Indices;
        std::uint64_t Triangles=0,SkirtTriangles=0;
        std::uint32_t TileCount=0;
    };

    template<typename T> bool Read(std::ifstream& File,T& Value)
    {
        File.read(reinterpret_cast<char*>(&Value),sizeof(Value));
        return static_cast<bool>(File);
    }

    bool LoadScene(const std::filesystem::path& Path,FScene& Scene,std::string& Error)
    {
        std::ifstream File(Path,std::ios::binary);
        if(!File){Error="cannot open mesh scene";return false;}
        char Magic[8]{};
        File.read(Magic,sizeof(Magic));
        if(!File||std::memcmp(Magic,"VFMSH001",8)!=0)
        {Error="mesh scene magic is invalid";return false;}
        std::uint32_t Version=0,ExpectedTiles=0;
        if(!Read(File,Version)||!Read(File,Scene.VoxelSizeMicrons)||!Read(File,ExpectedTiles)
            ||Version!=1||Scene.VoxelSizeMicrons==0||ExpectedTiles>1000000)
        {Error="mesh scene header is malformed";return false;}
        for(double& Value:Scene.CameraMeters)if(!Read(File,Value))
        {Error="mesh scene camera is truncated";return false;}
        if(!Read(File,Scene.YawDegrees)||!Read(File,Scene.PitchDegrees))
        {Error="mesh scene camera orientation is truncated";return false;}
        Scene.TileCount=ExpectedTiles;
        std::uint64_t TotalVertices=0,TotalIndices=0;
        for(std::uint32_t Tile=0;Tile<ExpectedTiles;++Tile)
        {
            std::int32_t Key[4]{};
            std::uint32_t SurfaceVertices=0,SkirtCount=0,VertexCount=0,IndexCount=0,MaskCount=0;
            for(auto& Value:Key)if(!Read(File,Value)){Error="tile key is truncated";return false;}
            if(!Read(File,SurfaceVertices)||!Read(File,SkirtCount)||!Read(File,VertexCount)
                ||!Read(File,IndexCount)||!Read(File,MaskCount)
                ||VertexCount>10000000||IndexCount>60000000||MaskCount>20000000
                ||IndexCount%3!=0||(MaskCount!=0&&MaskCount!=IndexCount/3))
            {Error="tile mesh header is malformed";return false;}
            if(TotalVertices+VertexCount>std::numeric_limits<std::uint32_t>::max()
                ||TotalIndices+IndexCount>std::numeric_limits<std::size_t>::max())
            {Error="mesh scene exceeds supported index space";return false;}
            const std::size_t Base=Scene.Vertices.size();
            Scene.Vertices.resize(Base+VertexCount);
            Scene.Indices.reserve(Scene.Indices.size()+IndexCount);
            for(std::size_t I=Base;I<Base+VertexCount;++I)
                for(float& Value:Scene.Vertices[I].Position)if(!Read(File,Value))
                {Error="tile positions are truncated";return false;}
            for(std::size_t I=Base;I<Base+VertexCount;++I)
            {
                for(float& Value:Scene.Vertices[I].Normal)if(!Read(File,Value))
                {Error="tile normals are truncated";return false;}
                Scene.Vertices[I].Height=Scene.Vertices[I].Position[2];
            }
            const std::size_t IndexOffset=Scene.Indices.size();
            Scene.Indices.resize(IndexOffset+IndexCount);
            if(IndexCount)File.read(reinterpret_cast<char*>(Scene.Indices.data()+IndexOffset),
                static_cast<std::streamsize>(sizeof(std::uint32_t)*IndexCount));
            std::vector<std::uint8_t> Masks(MaskCount);
            if(MaskCount)File.read(reinterpret_cast<char*>(Masks.data()),MaskCount);
            if(!File){Error="tile index or skirt data is truncated";return false;}
            for(std::size_t Index=IndexOffset;Index<Scene.Indices.size();++Index)
            {
                if(Scene.Indices[Index]>=VertexCount)
                {Error="tile index points outside its vertex array";return false;}
                Scene.Indices[Index]+=static_cast<std::uint32_t>(Base);
            }
            Scene.Triangles+=IndexCount/3;
            Scene.SkirtTriangles+=SkirtCount;
            TotalVertices+=VertexCount;TotalIndices+=IndexCount;
        }
        if(File.peek()!=std::char_traits<char>::eof())
        {Error="mesh scene has trailing bytes";return false;}
        if(Scene.Indices.empty()){Error="mesh scene contains no triangles";return false;}
        return true;
    }

    EGLDisplay MakeDisplay()
    {
        using TGetPlatformDisplay=EGLDisplay(*)(EGLenum,void*,const EGLint*);
        auto GetPlatformDisplay=reinterpret_cast<TGetPlatformDisplay>(
            eglGetProcAddress("eglGetPlatformDisplayEXT"));
#ifdef EGL_PLATFORM_SURFACELESS_MESA
        if(GetPlatformDisplay)
        {
            const EGLDisplay Display=GetPlatformDisplay(EGL_PLATFORM_SURFACELESS_MESA,
                EGL_DEFAULT_DISPLAY,nullptr);
            if(Display!=EGL_NO_DISPLAY)return Display;
        }
#endif
        return eglGetDisplay(EGL_DEFAULT_DISPLAY);
    }

    GLuint CompileShader(GLenum Type,const char* Source)
    {
        const GLuint Shader=glCreateShader(Type);
        glShaderSource(Shader,1,&Source,nullptr);
        glCompileShader(Shader);
        GLint Success=GL_FALSE;
        glGetShaderiv(Shader,GL_COMPILE_STATUS,&Success);
        if(Success==GL_FALSE)
        {
            GLint Length=0;glGetShaderiv(Shader,GL_INFO_LOG_LENGTH,&Length);
            std::string Log(static_cast<std::size_t>(std::max(Length,1)),0);
            glGetShaderInfoLog(Shader,Length,nullptr,Log.data());
            std::fprintf(stderr,"[OffscreenGL] shader compile failed: %s\n",Log.c_str());
            glDeleteShader(Shader);return 0;
        }
        return Shader;
    }

    GLuint MakeProgram()
    {
        static constexpr char VertexSource[]=R"GLSL(#version 300 es
precision highp float;
layout(location=0) in vec3 inCameraPosition;
layout(location=1) in vec3 inNormal;
layout(location=2) in float inHeight;
out vec3 vNormal;
out float vHeight;
uniform float uTanHalfHorizontal;
uniform float uTanHalfVertical;
uniform float uNear;
uniform float uFar;
void main()
{
    float z=max(inCameraPosition.z,0.0001);
    float clipZ=((uFar+uNear)/(uFar-uNear))*z-(2.0*uFar*uNear)/(uFar-uNear);
    gl_Position=vec4(inCameraPosition.x/uTanHalfHorizontal,
        inCameraPosition.y/uTanHalfVertical,clipZ,z);
    vNormal=normalize(inNormal);
    vHeight=inHeight;
})GLSL";
        static constexpr char FragmentSource[]=R"GLSL(#version 300 es
precision highp float;
in vec3 vNormal;
in float vHeight;
out vec4 outColor;
void main()
{
    vec3 n=normalize(vNormal);
    vec3 normalColor=abs(n)*0.78+vec3(0.12);
    float band=mod(floor((vHeight+10000.0)/80.0),2.0);
    vec3 bandColor=mix(vec3(0.24,0.31,0.37),vec3(0.39,0.42,0.35),band);
    outColor=vec4(mix(bandColor,normalColor,0.68),1.0);
})GLSL";
        const GLuint VS=CompileShader(GL_VERTEX_SHADER,VertexSource);
        const GLuint FS=CompileShader(GL_FRAGMENT_SHADER,FragmentSource);
        if(!VS||!FS)return 0;
        const GLuint Program=glCreateProgram();
        glAttachShader(Program,VS);glAttachShader(Program,FS);glLinkProgram(Program);
        glDeleteShader(VS);glDeleteShader(FS);
        GLint Success=GL_FALSE;glGetProgramiv(Program,GL_LINK_STATUS,&Success);
        if(Success==GL_FALSE)
        {
            GLint Length=0;glGetProgramiv(Program,GL_INFO_LOG_LENGTH,&Length);
            std::string Log(static_cast<std::size_t>(std::max(Length,1)),0);
            glGetProgramInfoLog(Program,Length,nullptr,Log.data());
            std::fprintf(stderr,"[OffscreenGL] program link failed: %s\n",Log.c_str());
            glDeleteProgram(Program);return 0;
        }
        return Program;
    }

    bool WritePng(const std::filesystem::path& Path,int Width,int Height,
                  const std::vector<std::uint8_t>& Pixels)
    {
        std::error_code Error;
        if(!Path.parent_path().empty())std::filesystem::create_directories(Path.parent_path(),Error);
        if(Error)return false;
        FILE* File=std::fopen(Path.c_str(),"wb");
        if(!File)return false;
        png_structp Png=png_create_write_struct(PNG_LIBPNG_VER_STRING,nullptr,nullptr,nullptr);
        png_infop Info=Png?png_create_info_struct(Png):nullptr;
        if(!Png||!Info){if(Png)png_destroy_write_struct(&Png,nullptr);std::fclose(File);return false;}
        if(setjmp(png_jmpbuf(Png)))
        {png_destroy_write_struct(&Png,&Info);std::fclose(File);return false;}
        png_init_io(Png,File);
        png_set_IHDR(Png,Info,static_cast<png_uint_32>(Width),static_cast<png_uint_32>(Height),
            8,PNG_COLOR_TYPE_RGBA,PNG_INTERLACE_NONE,PNG_COMPRESSION_TYPE_DEFAULT,PNG_FILTER_TYPE_DEFAULT);
        png_write_info(Png,Info);
        const std::size_t Stride=static_cast<std::size_t>(Width)*4;
        for(int Y=Height-1;Y>=0;--Y)
            png_write_row(Png,const_cast<png_bytep>(Pixels.data()+static_cast<std::size_t>(Y)*Stride));
        png_write_end(Png,nullptr);
        png_destroy_write_struct(&Png,&Info);
        const bool Good=std::fclose(File)==0;
        return Good;
    }

    bool Render(FScene& Scene,const std::filesystem::path& Output,int Width,int Height,
                double HorizontalFovDegrees)
    {
        EGLDisplay Display=MakeDisplay();
        EGLint Major=0,Minor=0;
        if(Display==EGL_NO_DISPLAY||eglInitialize(Display,&Major,&Minor)!=EGL_TRUE)
        {std::fprintf(stderr,"[OffscreenEGL] eglInitialize failed (0x%x)\n",eglGetError());return false;}
        if(eglBindAPI(EGL_OPENGL_ES_API)!=EGL_TRUE)
        {std::fprintf(stderr,"[OffscreenEGL] eglBindAPI failed (0x%x)\n",eglGetError());eglTerminate(Display);return false;}
        const EGLint ConfigAttributes[]={EGL_SURFACE_TYPE,EGL_PBUFFER_BIT,
            EGL_RENDERABLE_TYPE,EGL_OPENGL_ES3_BIT,EGL_RED_SIZE,8,EGL_GREEN_SIZE,8,
            EGL_BLUE_SIZE,8,EGL_ALPHA_SIZE,8,EGL_DEPTH_SIZE,24,EGL_NONE};
        EGLConfig Config{};EGLint ConfigCount=0;
        if(eglChooseConfig(Display,ConfigAttributes,&Config,1,&ConfigCount)!=EGL_TRUE||ConfigCount==0)
        {std::fprintf(stderr,"[OffscreenEGL] no GLES3 pbuffer config (0x%x)\n",eglGetError());eglTerminate(Display);return false;}
        const EGLint SurfaceAttributes[]={EGL_WIDTH,Width,EGL_HEIGHT,Height,EGL_NONE};
        EGLSurface Surface=eglCreatePbufferSurface(Display,Config,SurfaceAttributes);
        const EGLint ContextAttributes[]={EGL_CONTEXT_CLIENT_VERSION,3,EGL_NONE};
        EGLContext Context=eglCreateContext(Display,Config,EGL_NO_CONTEXT,ContextAttributes);
        if(Surface==EGL_NO_SURFACE||Context==EGL_NO_CONTEXT
            ||eglMakeCurrent(Display,Surface,Surface,Context)!=EGL_TRUE)
        {std::fprintf(stderr,"[OffscreenEGL] context/surface creation failed (0x%x)\n",eglGetError());
            if(Context!=EGL_NO_CONTEXT)eglDestroyContext(Display,Context);
            if(Surface!=EGL_NO_SURFACE)eglDestroySurface(Display,Surface);
            eglTerminate(Display);return false;}

        const GLubyte* Renderer=glGetString(GL_RENDERER);
        const GLubyte* GlVersion=glGetString(GL_VERSION);
        std::printf("[OffscreenEGL] renderer=%s version=%s pbuffer=%dx%d\n",
            Renderer?reinterpret_cast<const char*>(Renderer):"unknown",
            GlVersion?reinterpret_cast<const char*>(GlVersion):"unknown",Width,Height);
        GLuint ColorTexture=0,DepthBuffer=0,Framebuffer=0;
        glGenTextures(1,&ColorTexture);glBindTexture(GL_TEXTURE_2D,ColorTexture);
        glTexImage2D(GL_TEXTURE_2D,0,GL_RGBA8,Width,Height,0,GL_RGBA,GL_UNSIGNED_BYTE,nullptr);
        glTexParameteri(GL_TEXTURE_2D,GL_TEXTURE_MIN_FILTER,GL_NEAREST);
        glTexParameteri(GL_TEXTURE_2D,GL_TEXTURE_MAG_FILTER,GL_NEAREST);
        glGenRenderbuffers(1,&DepthBuffer);glBindRenderbuffer(GL_RENDERBUFFER,DepthBuffer);
        glRenderbufferStorage(GL_RENDERBUFFER,GL_DEPTH_COMPONENT24,Width,Height);
        glGenFramebuffers(1,&Framebuffer);glBindFramebuffer(GL_FRAMEBUFFER,Framebuffer);
        glFramebufferTexture2D(GL_FRAMEBUFFER,GL_COLOR_ATTACHMENT0,GL_TEXTURE_2D,ColorTexture,0);
        glFramebufferRenderbuffer(GL_FRAMEBUFFER,GL_DEPTH_ATTACHMENT,GL_RENDERBUFFER,DepthBuffer);
        if(glCheckFramebufferStatus(GL_FRAMEBUFFER)!=GL_FRAMEBUFFER_COMPLETE)
        {std::fprintf(stderr,"[OffscreenGL] framebuffer is incomplete\n");return false;}

        constexpr double Pi=3.14159265358979323846;
        const double Yaw=Scene.YawDegrees*Pi/180.0,Pitch=Scene.PitchDegrees*Pi/180.0;
        const std::array<double,3> Forward={std::sin(Yaw)*std::cos(Pitch),
            std::cos(Yaw)*std::cos(Pitch),std::sin(Pitch)};
        const std::array<double,3> Right={std::cos(Yaw),-std::sin(Yaw),0.0};
        const std::array<double,3> Up={Right[1]*Forward[2]-Right[2]*Forward[1],
            Right[2]*Forward[0]-Right[0]*Forward[2],Right[0]*Forward[1]-Right[1]*Forward[0]};
        const double TanHalfHorizontal=std::tan(HorizontalFovDegrees*Pi/360.0);
        const double Aspect=static_cast<double>(Width)/Height;
        const double TanHalfVertical=TanHalfHorizontal/Aspect;
        const double Near=0.1,Far=20000.0;
        for(FVertex& Vertex:Scene.Vertices)
        {
            const double DX=Vertex.Position[0]-Scene.CameraMeters[0];
            const double DY=Vertex.Position[1]-Scene.CameraMeters[1];
            const double DZ=Vertex.Position[2]-Scene.CameraMeters[2];
            const float CameraX=static_cast<float>(DX*Right[0]+DY*Right[1]+DZ*Right[2]);
            const float CameraY=static_cast<float>(DX*Up[0]+DY*Up[1]+DZ*Up[2]);
            const float CameraZ=static_cast<float>(DX*Forward[0]+DY*Forward[1]+DZ*Forward[2]);
            Vertex.Position[0]=CameraX;Vertex.Position[1]=CameraY;Vertex.Position[2]=CameraZ;
        }
        const GLuint Program=MakeProgram();
        if(!Program)return false;
        GLuint Vao=0,Vbo=0,Ebo=0;
        glGenVertexArrays(1,&Vao);glBindVertexArray(Vao);
        glGenBuffers(1,&Vbo);glBindBuffer(GL_ARRAY_BUFFER,Vbo);
        glBufferData(GL_ARRAY_BUFFER,static_cast<GLsizeiptr>(Scene.Vertices.size()*sizeof(FVertex)),
            Scene.Vertices.data(),GL_STATIC_DRAW);
        glGenBuffers(1,&Ebo);glBindBuffer(GL_ELEMENT_ARRAY_BUFFER,Ebo);
        glBufferData(GL_ELEMENT_ARRAY_BUFFER,static_cast<GLsizeiptr>(Scene.Indices.size()*sizeof(std::uint32_t)),
            Scene.Indices.data(),GL_STATIC_DRAW);
        glEnableVertexAttribArray(0);glVertexAttribPointer(0,3,GL_FLOAT,GL_FALSE,sizeof(FVertex),
            reinterpret_cast<void*>(offsetof(FVertex,Position)));
        glEnableVertexAttribArray(1);glVertexAttribPointer(1,3,GL_FLOAT,GL_FALSE,sizeof(FVertex),
            reinterpret_cast<void*>(offsetof(FVertex,Normal)));
        glEnableVertexAttribArray(2);glVertexAttribPointer(2,1,GL_FLOAT,GL_FALSE,sizeof(FVertex),
            reinterpret_cast<void*>(offsetof(FVertex,Height)));
        glViewport(0,0,Width,Height);glDisable(GL_CULL_FACE);glEnable(GL_DEPTH_TEST);
        glDepthFunc(GL_LEQUAL);glClearColor(0.035f,0.055f,0.075f,1.0f);
        glClearDepthf(1.0f);glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT);
        glUseProgram(Program);
        glUniform1f(glGetUniformLocation(Program,"uTanHalfHorizontal"),static_cast<float>(TanHalfHorizontal));
        glUniform1f(glGetUniformLocation(Program,"uTanHalfVertical"),static_cast<float>(TanHalfVertical));
        glUniform1f(glGetUniformLocation(Program,"uNear"),static_cast<float>(Near));
        glUniform1f(glGetUniformLocation(Program,"uFar"),static_cast<float>(Far));
        glDrawElements(GL_TRIANGLES,static_cast<GLsizei>(Scene.Indices.size()),GL_UNSIGNED_INT,nullptr);
        glFinish();
        std::vector<std::uint8_t> Pixels(static_cast<std::size_t>(Width)*Height*4);
        glReadPixels(0,0,Width,Height,GL_RGBA,GL_UNSIGNED_BYTE,Pixels.data());
        const GLenum GlError=glGetError();
        const bool PngGood=GlError==GL_NO_ERROR&&WritePng(Output,Width,Height,Pixels);
        if(!PngGood)std::fprintf(stderr,"[OffscreenGL] render or PNG write failed (GL=0x%x)\n",GlError);
        else
        {
            std::error_code SizeError;
            const auto Bytes=std::filesystem::file_size(Output,SizeError);
            std::printf("[OffscreenCapture] triangles=%llu tiles=%u skirts=%llu png_bytes=%llu status=PASS path=%s\n",
                static_cast<unsigned long long>(Scene.Triangles),Scene.TileCount,
                static_cast<unsigned long long>(Scene.SkirtTriangles),
                static_cast<unsigned long long>(SizeError?0:Bytes),Output.string().c_str());
        }
        glDeleteProgram(Program);glDeleteBuffers(1,&Ebo);glDeleteBuffers(1,&Vbo);
        glDeleteVertexArrays(1,&Vao);glDeleteFramebuffers(1,&Framebuffer);
        glDeleteRenderbuffers(1,&DepthBuffer);glDeleteTextures(1,&ColorTexture);
        eglMakeCurrent(Display,EGL_NO_SURFACE,EGL_NO_SURFACE,EGL_NO_CONTEXT);
        eglDestroyContext(Display,Context);eglDestroySurface(Display,Surface);eglTerminate(Display);
        return PngGood;
    }
}

int main(int argc,char** argv)
{
    if(argc<3)
    {
        std::fprintf(stderr,"Usage: %s scene.vfmesh output.png [--size=WIDTHxHEIGHT] [--fov-deg=90] [--camera-m=X,Y,Z] [--yaw-deg=N] [--pitch-deg=N]\n",argv[0]);
        return 2;
    }
    int Width=2560,Height=1440;double Fov=90.0;
    bool HasYaw=false,HasPitch=false;double Yaw=0.0,Pitch=0.0;
    bool HasCamera=false;std::array<double,3> Camera{};
    for(int I=3;I<argc;++I)
    {
        const std::string_view Arg(argv[I]);
        if(Arg.starts_with("--size="))
        {
            const std::string Value(Arg.substr(7));
            if(std::sscanf(Value.c_str(),"%dx%d",&Width,&Height)!=2||Width<=0||Height<=0)
            {std::fprintf(stderr,"Invalid image size\n");return 2;}
        }
        else if(Arg.starts_with("--fov-deg="))
        {
            char* End=nullptr;const std::string Value(Arg.substr(10));Fov=std::strtod(Value.c_str(),&End);
            if(End==Value.c_str()||*End!='\0'||Fov<=1.0||Fov>=179.0)
            {std::fprintf(stderr,"Invalid horizontal FOV\n");return 2;}
        }
        else if(Arg.starts_with("--camera-m="))
        {
            const std::string Value(Arg.substr(11));
            if(std::sscanf(Value.c_str(),"%lf,%lf,%lf",&Camera[0],&Camera[1],&Camera[2])!=3
                ||!std::isfinite(Camera[0])||!std::isfinite(Camera[1])||!std::isfinite(Camera[2]))
            {std::fprintf(stderr,"Invalid camera position\n");return 2;}
            HasCamera=true;
        }
        else if(Arg.starts_with("--yaw-deg=")||Arg.starts_with("--pitch-deg="))
        {
            const bool IsYaw=Arg.starts_with("--yaw-deg=");
            const std::size_t Prefix=IsYaw?10:12;
            const std::string Value(Arg.substr(Prefix));char* End=nullptr;
            const double Parsed=std::strtod(Value.c_str(),&End);
            if(End==Value.c_str()||*End!='\0'||!std::isfinite(Parsed))
            {std::fprintf(stderr,"Invalid camera angle\n");return 2;}
            if(IsYaw){Yaw=Parsed;HasYaw=true;}else{Pitch=Parsed;HasPitch=true;}
        }
        else {std::fprintf(stderr,"Unknown option: %.*s\n",static_cast<int>(Arg.size()),Arg.data());return 2;}
    }
    FScene Scene;std::string Error;
    if(!LoadScene(argv[1],Scene,Error))
    {std::fprintf(stderr,"[OffscreenCapture] FAIL: %s\n",Error.c_str());return 2;}
    if(HasCamera)Scene.CameraMeters=Camera;
    if(HasYaw)Scene.YawDegrees=Yaw;
    if(HasPitch)Scene.PitchDegrees=Pitch;
    std::printf("[OffscreenScene] voxels_per_meter=%.6f camera_m=(%.3f,%.3f,%.3f) yaw=%.2f pitch=%.2f vertices=%zu triangles=%llu\n",
        1000000.0/Scene.VoxelSizeMicrons,Scene.CameraMeters[0],Scene.CameraMeters[1],Scene.CameraMeters[2],
        Scene.YawDegrees,Scene.PitchDegrees,Scene.Vertices.size(),
        static_cast<unsigned long long>(Scene.Triangles));
    return Render(Scene,argv[2],Width,Height,Fov)?0:3;
}