/**********************************************************************************************************

Copyright 2021. Lukasz Czyz

Permission is hereby granted, free of charge, to any person obtaining a copy of this file, to deal with the file without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the file, and to permit persons to whom the file is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

THE ABOVE COPYRIGHT NOTICE APPLIES TO THIS SAMPLE FILE ONLY. THE UVPACKMASTER SDK IS DISTRIBUTED UNDER A DIFFERENT EULA.

***********************************************************************************************************/

// The main UVPackmaster SDK header
#include <uvpCore.hpp>

#ifdef __linux__
    // For getopt
    #include <unistd.h>
#elif _WIN32
    // Provide an implementation of getopt on your own if you want to compile this program on Windows.
    #include "getopt.hpp"
#else
    #error "Unsupported platform"
#endif

// Download this header from https://github.com/datenwolf/linmath.h if you want to compile this program.
#include "linmath.h"

// FBX SDK is not distributed together with this program, get the SDK files on your own if you want to compile it.
#include <fbxsdk.h>

#include <iostream>
#include <unordered_map>
#include <list>
#include <vector>
#include <string>
#include <cstring>
#include <array>


using namespace uvpcore;


typedef std::array<UvpMessageT*, static_cast<int>(UvpMessageT::MESSAGE_CODE::VALUE_COUNT)> UvpMessageArrayT;
void opExecutorMessageHandler(void* m_pMessageHandlerData, UvpMessageT *pMsg);

// Wrapper class simplifying execution of UVP operations.
class UvpOpExecutorT
{
private:
    friend void opExecutorMessageHandler(void* m_pMessageHandlerData, UvpMessageT *pMsg);

    std::list<UvpMessageT*> m_ReceivedMessages;
    UvpMessageArrayT m_LastMessagePerCode;

    bool m_DebugMode;

    void destroyMessages()
    {       
        // The application becomes the owner of UVP messages after receiving it,
        // so we have to make sure they are eventually deallocated by calling
        // the destory method on them (do not use the delete operator).
        for (UvpMessageT *pMsg : m_ReceivedMessages)
        {
            pMsg->destroy();
        }
        m_ReceivedMessages.clear();
    }

    void reset()
    {
        destroyMessages();
        m_LastMessagePerCode = { nullptr };
    }

    void handleMessage(UvpMessageT *pMsg)
    {
        // This method is called every time the packer sends a message to the application.
        // We need to handle the message properly.

        if (pMsg->m_Code == UvpMessageT::MESSAGE_CODE::PROGRESS_REPORT)
        {
            UvpProgressReportMessageT *pReportProgressMsg = static_cast<UvpProgressReportMessageT*>(pMsg);

            std::cout << "[UVP PROGRESS REPORT] Phase: " << static_cast<int>(pReportProgressMsg->m_PackingPhase);
            for (int i = 0; i < pReportProgressMsg->m_ProgressSize; i++)
            {
                std::cout << ", Progress[" << i << "]: " << pReportProgressMsg->m_ProgressArray[i];
            }
            std::cout << "\n";
        }
        else if (pMsg->m_Code == UvpMessageT::MESSAGE_CODE::BENCHMARK)
        {
            UvpBenchmarkMessageT *pBenchmarkMsg = static_cast<UvpBenchmarkMessageT*>(pMsg);

            std::cout << "[UVP BENCHMARK] Device name: " << pBenchmarkMsg->m_DeviceName.c_str() << ", Total packing time (ms): " <<
                pBenchmarkMsg->m_TotalPackTimeMs << ", Average packing time (ms): " << pBenchmarkMsg->m_AvgPackTimeMs << "\n";
        }

        m_LastMessagePerCode[static_cast<int>(pMsg->m_Code)] = pMsg;
        m_ReceivedMessages.push_back(pMsg);
    }


public:
    UvpOpExecutorT(bool debugMode) :
        m_DebugMode(debugMode)
    {}

    ~UvpOpExecutorT()
    {
        destroyMessages();
    }

    UVP_ERRORCODE execute(UvpOperationInputT &uvpInput)
    {
        reset();

        uvpInput.m_pMessageHandler = opExecutorMessageHandler;
        uvpInput.m_pMessageHandlerData = this;

        if (m_DebugMode)
        {
            // Check whether the application configurated the operation input properly.
            // WARNING: this operation is time consuming (in particular it iterates over all UV data),
            // that is why it should only be executed when debugging the application. It should
            // never be used in production.
            const char *pValidationResult = uvpInput.validate();

            if (pValidationResult)
            {
                throw std::runtime_error("Operation input validation failed: " + std::string(pValidationResult));
            }
        }

        UvpOperationT uvpOp(uvpInput);
        // Execute the actual operation - this call will block the current thread
        // until the operation is finished.
        UVP_ERRORCODE retCode = uvpOp.entry();

        return retCode;
    }

    UvpMessageT* getLastMessage(UvpMessageT::MESSAGE_CODE code)
    {
        return m_LastMessagePerCode[static_cast<int>(code)];
    }
};

void opExecutorMessageHandler(void* m_pMessageHandlerData, UvpMessageT *pMsg)
{
    // This handler is called every time the packer sends a message to the application.
    // Simply pass the message to the underlaying executor object.
    reinterpret_cast<UvpOpExecutorT*>(m_pMessageHandlerData)->handleMessage(pMsg);
}


void islandSolutionToMatrix(const UvpIslandPackSolutionT &islandSolution, mat4x4 &mat)
{
    // Generate matrix used to transform UVs of the given island in order to apply
    // a packing result
    mat4x4_identity(mat);
    mat4x4_translate_in_place(mat, islandSolution.m_PostScaleOffset[0], islandSolution.m_PostScaleOffset[1], 0.0);
    mat4x4_scale_aniso(mat, mat, 1.0 / islandSolution.m_Scale, 1.0 / islandSolution.m_Scale, 1.0);
    mat4x4_translate_in_place(mat, islandSolution.m_Offset[0], islandSolution.m_Offset[1], 0.0);
    mat4x4_translate_in_place(mat, islandSolution.m_Pivot[0], islandSolution.m_Pivot[1], 0.0);
    mat4x4_rotate_Z(mat, mat, islandSolution.m_Angle);
    mat4x4_translate_in_place(mat, -islandSolution.m_Pivot[0], -islandSolution.m_Pivot[1], 0.0);
    mat4x4_scale_aniso(mat, mat, islandSolution.m_PreScale, islandSolution.m_PreScale, 1.0);
}


class FbxUvWrapper
{
private:
    FbxManager *m_pSdkManager = nullptr;
    FbxScene *m_pScene = nullptr;
    FbxMesh *m_pMesh = nullptr;
    FbxGeometryElementUV* m_pUvElement = nullptr;

    int m_PolyVertexCount;

    std::vector<UvVertT> m_VertArray;
    std::vector<UvFaceT> m_FaceArray;

public:
    FbxUvWrapper(const char *pFbxFilepath, const char *pObjName, UvDataT &uvData)
    {
        m_pSdkManager = FbxManager::Create();

        FbxIOSettings *ios = FbxIOSettings::Create(m_pSdkManager, IOSROOT);
        m_pSdkManager->SetIOSettings(ios);

        FbxImporter *pImporter = FbxImporter::Create(m_pSdkManager, "");
        bool importStatus = pImporter->Initialize(pFbxFilepath, -1, m_pSdkManager->GetIOSettings());

        if (!importStatus)
        {
            throw std::runtime_error("Could not import the fbx file");
        }

        m_pScene = FbxScene::Create(m_pSdkManager,"");

        pImporter->Import(m_pScene);
        pImporter->Destroy();

        FbxNode *pObjFbxNode = nullptr;

        for (int i = 0; i < m_pScene->GetNodeCount(); i++)
        {
            FbxNode *pNode = m_pScene->GetNode(i);
            
            if (std::strcmp(pObjName, pNode->GetName()) == 0)
            {
                pObjFbxNode = pNode;
                break;
            }
        }

        if (pObjFbxNode == nullptr)
        {
            throw std::runtime_error("Object not found in the fbx file");
        }

        FbxNodeAttribute *pObjNodeAttr = pObjFbxNode->GetNodeAttribute();

        if (pObjNodeAttr == nullptr)
        {
            throw std::runtime_error("Object node does not have a node attribute");
        }

        if (pObjNodeAttr->GetAttributeType() != FbxNodeAttribute::EType::eMesh)
        {
            throw std::runtime_error("Object node attribute is not a mesh");
        }

        m_pMesh = static_cast<FbxMesh*>(pObjNodeAttr);

        FbxStringList uvSetNames;
        m_pMesh->GetUVSetNames(uvSetNames);

        if (uvSetNames.GetCount() == 0)
        {
            throw std::runtime_error("Object mesh does not have UV map");
        }

        const char *pUvSetName = uvSetNames[0];

        m_pUvElement = m_pMesh->GetElementUV(pUvSetName);
        if(!m_pUvElement)
        {
            throw std::runtime_error("Required UV element not found");
        }

        if (m_pUvElement->GetMappingMode() != FbxGeometryElement::eByPolygonVertex &&
            m_pUvElement->GetMappingMode() != FbxGeometryElement::eByControlPoint)
        {
            throw std::runtime_error("Unsupported mapping mode");
        }


        m_PolyVertexCount = 0;

        // By using UvVertHashT and UvVertEqualT we identify UV vertices having
        // the same UV coordinates and corresponding to the same 3D vertex.
        // By using this approach we can make sure no UV face contains duplicated vertices i.e.
        // the same UV vertex index is not present in the vertex array of a single face twice.
        // This is required as duplicated vertices are not allowed by the packer.
        std::unordered_map<UvVertT, int, UvVertHashT, UvVertEqualT> vertPointerMap;

        for (int faceIdx = 0; faceIdx < m_pMesh->GetPolygonCount(); faceIdx++)
        {
            int faceSize = m_pMesh->GetPolygonSize(faceIdx);
            m_PolyVertexCount += faceSize;

            m_FaceArray.emplace_back(faceIdx);
            UvFaceT &face = m_FaceArray.back();

            auto &faceVerts = face.m_Verts;
            faceVerts.reserve(faceSize);

            for (int vertIdx = 0; vertIdx < faceSize; vertIdx++)
            {
                FbxVector2 uvs;
                bool unmapped;

                m_pMesh->GetPolygonVertexUV(
                    faceIdx,
                    vertIdx,
                    pUvSetName,
                    uvs,
                    unmapped);

                if (unmapped)
                {
                    throw std::runtime_error("Object mesh has unmapped vertices");
                }

                UvVertT uvVert;

                uvVert.m_UvCoords[0] = uvs[0];
                uvVert.m_UvCoords[1] = uvs[1];
                uvVert.m_ControlId = m_pMesh->GetPolygonVertex(faceIdx, vertIdx);

                auto it = vertPointerMap.find(uvVert);

                size_t newVertIdx;
                if (it == vertPointerMap.end())
                {
                    // It is the first time we see such a vertex - add it to the array
                    newVertIdx = m_VertArray.size();
                    vertPointerMap[uvVert] = newVertIdx;
                    m_VertArray.emplace_back(uvVert);
                }
                else
                {
                    // Such a vertex has already been added to the array.
                    newVertIdx = (*it).second;
                }

                faceVerts.pushBack(newVertIdx);
            }
        }

        // Initialize UvData which will be passed to the packer with the generated arrays
        if (m_FaceArray.size() > 0)
        {
            uvData.m_FaceCount = m_FaceArray.size();
            uvData.m_pFaceArray = m_FaceArray.data();
        }

        if (m_VertArray.size() > 0)
        {
            uvData.m_VertCount = m_VertArray.size();
            uvData.m_pVertArray = m_VertArray.data();
        }
    }

    ~FbxUvWrapper()
    {
        // Calling Destory on the Sdk manager will destory all related SDK objects.
        if (m_pSdkManager)
        {
            m_pSdkManager->Destroy();
        }
    }

    void applyPackResult(const UvpIslandsMessageT *pIslandsMsg, const UvpPackSolutionMessageT *pPackSolutionMsg)
    {
        // Array of UV vertices transformed by the packing result.
        std::vector<FbxVector2> transformedUvs(m_VertArray.size());

        // Initially copy the original UV coordinates.
        for (int i = 0; i < m_VertArray.size(); i++)
        {
            const auto &origVert = m_VertArray[i];
            transformedUvs[i] = FbxVector2(origVert.m_UvCoords[0], origVert.m_UvCoords[1]);
        }

        // Transform UV coordinates accordingly
        const auto &islands = pIslandsMsg->m_Islands;
        for (const UvpIslandPackSolutionT &islandSolution : pPackSolutionMsg->m_IslandSolutions)
        {
            const IdxArrayT &island = islands[islandSolution.m_IslandIdx];
            mat4x4 solutionMatrix;
            islandSolutionToMatrix(islandSolution, solutionMatrix);

            for (int faceId : island)
            {
                const UvFaceT &face = m_FaceArray[faceId];

                for (int vertIdx : face.m_Verts)
                {
                    const UvVertT &origVert = m_VertArray[vertIdx];
                    vec4 inputUv = { origVert.m_UvCoords[0], origVert.m_UvCoords[1], 0.0, 1.0 };
                    vec4 transformedUv;

                    mat4x4_mul_vec4(transformedUv, solutionMatrix, inputUv);
                    transformedUvs[vertIdx] = FbxVector2(transformedUv[0] / transformedUv[3], transformedUv[1] / transformedUv[3]);
                }
            }
        }

        const bool useIndex = m_pUvElement->GetReferenceMode() != FbxGeometryElement::eDirect;

        if (useIndex)
        {
            m_pUvElement->GetIndexArray().Resize(m_PolyVertexCount);
            m_pUvElement->GetDirectArray().Resize(transformedUvs.size());
        }
        else
        {
            m_pUvElement->GetDirectArray().Resize(m_PolyVertexCount);
        }

        int polyIdxCounter = 0;

        for (const UvFaceT &uvFace : m_FaceArray)
        {
            for (const int vertIdx : uvFace.m_Verts)
            {
                if (useIndex)
                {
                    m_pUvElement->GetIndexArray().SetAt(polyIdxCounter, vertIdx);
                    m_pUvElement->GetDirectArray().SetAt(vertIdx, transformedUvs[vertIdx]);
                }
                else
                {
                    m_pUvElement->GetDirectArray().SetAt(polyIdxCounter, transformedUvs[vertIdx]);
                }

                polyIdxCounter++;
            }
        }
    }

    void saveToFile(const char* pOutFilename)
    {
        FbxExporter* pExporter = FbxExporter::Create(m_pSdkManager, "");

        int format = m_pSdkManager->GetIOPluginRegistry()->GetNativeWriterFormat();
        if(!pExporter->Initialize(pOutFilename, format, m_pSdkManager->GetIOSettings()))
        {
            throw std::runtime_error("Could not initialize exporter");
        }

        if (!pExporter->Export(m_pScene))
        {
            throw std::runtime_error("Could not export scene");
        }

        pExporter->Destroy();
    }
};


int main(int argc, char * const argv[])
{
    int errCode = 0;

    try
    {
        char optChar;
        const char *pFbxFilePath = nullptr;
        const char *pOutFilePath = nullptr;
        const char *pMeshName = nullptr;
        bool debugMode = false;

        UvpOperationInputT uvpInput;

        // Set some parameters by default
        uvpInput.m_pDeviceId = "cpu";
        uvpInput.m_RenderResult = true;
        uvpInput.m_RenderInvalidIslands = true;
        uvpInput.m_RealtimeSolution = true;
        uvpInput.m_Benchmark = true;
        uvpInput.m_Opcode = UVP_OPCODE::PACK;

        // Set some parameters from the command line arguments. In order to execute a simple packing scenario
        // you only need to provide '-f' and '-m' arguments in command line.
        while ((optChar = getopt(argc, argv, "d:f:o:i:m:M:N:y:W:OF:r:Rg:a:t:Lp:Dh:j:HY:q:V:C:wxsI:l:QvB:U:")) != -1)
        {
            switch (optChar) {
            case 'd':
                uvpInput.m_pDeviceId = optarg;
                break;
            case 'f':
                pFbxFilePath = optarg;
                break;
            case 'o':
                pOutFilePath = optarg;
                break;
            case 'm':
                pMeshName = optarg;
                break;
            case 'i':
                uvpInput.m_Precision = std::stoi(optarg);
                break;
            case 'B':
                uvpInput.m_Margin = std::stof(optarg);
                break;
            case 'L':
                uvpInput.m_NormalizeIslands = true;
                break;
            case 'M':
                uvpInput.m_PixelMargin = std::stof(optarg);
                break;
            case 'N':
                uvpInput.m_PixelPadding = std::stof(optarg);
                break;
            case 'y':
                uvpInput.m_PixelMarginTextureSize = std::stoi(optarg);
                break;
            case 'W':
                uvpInput.m_PixelMarginMethod = static_cast<UVP_PIXEL_MARGIN_METHOD>(std::stoi(optarg));
                break;
            case 'O':
                uvpInput.m_FixedScale = true;
                break;
            case 'F':
                uvpInput.m_FixedScaleStrategy = static_cast<UVP_FIXED_SCALE_STRATEGY>(std::stoi(optarg));
                break;
            case 'r':
                uvpInput.m_RotationStep = std::stoi(optarg);
                break;
            case 'R':
                uvpInput.m_IslandRotationStepEnable = true;
                break;
            case 'w':
                uvpInput.m_PrerotDisable = true;
                break;
            case 't':
                uvpInput.m_ThreadCount = std::stoi(optarg);
                break;
            case 'h':
                uvpInput.m_HeuristicSearchTime = std::stoi(optarg);
                break;
            case 'j':
                uvpInput.m_HeuristicMaxWaitTime = std::stoi(optarg);
                break;
            case 'Y':
                uvpInput.m_PixelMarginAdjustTime = std::stoi(optarg);
                break;
            case 'H':
                uvpInput.m_AdvancedHeuristic = true;
                break;
            case 'q':
                uvpInput.m_PackRatio = std::stof(optarg);
                break;
            case 'x':
                uvpInput.m_PackToOthers = true;
                break;
            case 's':
                uvpInput.m_ProcessUnselected = true;
                break;
            case 'I':
                uvpInput.m_SimilarityThreshold = std::stof(optarg);
                break;
            case 'g':
                uvpInput.m_PackingMode = static_cast<UVP_PACKING_MODE>(std::stoi(optarg));
                break;
            case 'V':
                uvpInput.m_TileCount = std::stoi(optarg);
                break;
            case 'C':
                uvpInput.m_TilesInRow = std::stoi(optarg);
                break;
            case 'a':
                uvpInput.m_GroupingMethod = static_cast<UVP_GROUPING_METHOD>(std::stoi(optarg));
                break;
            case 'U':
                uvpInput.m_GroupingCompactness = std::stof(optarg);
                break;
            case 'l':
                uvpInput.m_LockOverlappingMode = static_cast<UVP_LOCK_OVERLAPPING_MODE>(std::stoi(optarg));;
                break;
            case 'Q':
                uvpInput.m_LockGroupsEnable = true;
                break;
            case 'v':
                uvpInput.m_PreValidate = true;
                break;
            case 'D':
                debugMode = true;
                break;

            default:
                throw std::runtime_error("Unsupported option");
            }
        }

        UvpOpExecutorT opExecutor(debugMode);

        UvpOperationInputT reportVersionInput;
        reportVersionInput.m_Opcode = UVP_OPCODE::REPORT_VERSION;

        if (opExecutor.execute(reportVersionInput) != UVP_ERRORCODE::SUCCESS)
        {
            throw std::runtime_error("Report version op failed");
        }

        UvpMessageT *pVersionMessageBase = opExecutor.getLastMessage(UvpMessageT::MESSAGE_CODE::VERSION);
        if (!pVersionMessageBase)
        {
            throw std::runtime_error("Expected Version message not found");
        }

        UvpVersionMessageT *pVersionMessage = static_cast<UvpVersionMessageT*>(pVersionMessageBase);

        std::cout << "UVP core info:\n";
        std::cout << "Version: " << pVersionMessage->m_VersionMajor << "." << pVersionMessage->m_VersionMinor << "." << pVersionMessage->m_VersionPatch << "\n";
        std::cout << "Available packing devices in the system:\n";

        for (auto &devDesc : pVersionMessage->m_DeviceDescArray)
        {
            std::cout << "ID: " << devDesc.m_Id.c_str() << ", NAME: " << devDesc.m_Name.c_str() << ", SUPPORTED: " << devDesc.m_Supported << "\n";
        }


        if (!pFbxFilePath)
        {
            throw std::runtime_error("Fbx file path must be provided in order to perform a packing operation, aborting");
        }

        if (!pMeshName)
        {
            throw std::runtime_error("Fbx object name must be provided in order to perform a packing operation, aborting");
        }

        FbxUvWrapper fbxWrapper(pFbxFilePath, pMeshName, uvpInput.m_UvData);

        if (opExecutor.execute(uvpInput) != UVP_ERRORCODE::SUCCESS)
        {
            throw std::runtime_error("Packing operation failed");
        }

        if (pOutFilePath)
        {
            if (!opExecutor.getLastMessage(UvpMessageT::MESSAGE_CODE::ISLANDS) || !opExecutor.getLastMessage(UvpMessageT::MESSAGE_CODE::PACK_SOLUTION))
            {
                throw std::runtime_error("Expected UVP messages not found");
            }

            const UvpIslandsMessageT *pIslandsMsg = static_cast<const UvpIslandsMessageT*>(opExecutor.getLastMessage(UvpMessageT::MESSAGE_CODE::ISLANDS));
            const UvpPackSolutionMessageT *pPackSolutionMsg = static_cast<const UvpPackSolutionMessageT*>(opExecutor.getLastMessage(UvpMessageT::MESSAGE_CODE::PACK_SOLUTION));

            fbxWrapper.applyPackResult(pIslandsMsg, pPackSolutionMsg);
            fbxWrapper.saveToFile(pOutFilePath);
        }
    }
    catch (const std::exception &ex)
    {
        std::cerr << "Error: " << ex.what() << "\n";
        errCode = 1;
    }
    catch (...)
    {
        std::cerr << "Unexpected error\n";
        errCode = 1;
    }

    return errCode;
}