#include "PfxHair.h"
#include "utils.h"

#include <maya/MFnPfxGeometry.h>
#include <maya/MRenderLineArray.h>
#include <maya/MSelectionList.h>
#include <maya/MFnDagNode.h>
#include <maya/MItDependencyGraph.h>
#include <maya/MNodeClass.h>
#include <maya/MFn.h>
#include <maya/MDagPath.h>
#include <maya/MMeshIntersector.h>
#include <maya/MGlobal.h>
#include <maya/MFnMesh.h>
#include <ai_msg.h>

PfxHairHelper::PfxHairHelper(){};

PfxHairHelper::PfxHairHelper(MObject &pfxHair)
    : m_pfxHair(pfxHair)
{
    m_hairSystem = GetHairSystem();
}

PfxHairHelper::PfxHairHelper(string pfxHair)
{
    GetMayaObject(pfxHair, m_pfxHair);
    m_hairSystem = GetHairSystem();
}

PfxHairHelper::~PfxHairHelper(){};

MObject PfxHairHelper::GetHairSystem()
{
    MObject hairSystem;
    MFnPfxGeometry fnPfxHair(m_pfxHair);
    MObjectArray hairSystems = GetTypeNodes(fnPfxHair.findPlug("renderHairs"),
                                            MItDependencyGraph::kUpstream,
                                            MFn::kHairSystem);
    if (hairSystems.length())
        hairSystem = hairSystems[0];
    return hairSystem;
}

MRenderLineArray PfxHairHelper::GetHairData(MStatus *status)
{
    MFnPfxGeometry fnPfxNode(m_pfxHair);
    AiMsgInfo("[%s] : PfxNode : %s", kTranslatorName, fnPfxNode.fullPathName().asChar());
    MDagPath pfxPath = fnPfxNode.dagPath();
    MRenderLineArray renderLineArray, leafLines, flowerLines;
    fnPfxNode.getLineData(renderLineArray, leafLines, flowerLines, true, false, false, false, false, false, false, false, true);
    return renderLineArray;
}

int PfxHairHelper::GetHairCount()
{
    MStatus status;
    MRenderLineArray renderLineArray = GetHairData(&status);
    if (status == MS::kFailure)
        return 0;
    return renderLineArray.length();
}

MIntArray PfxHairHelper::GetHairBaseSegement(MStatus *status)
{
    MIntArray segmentArray;
    MRenderLineArray renderLineArray = GetHairData();
    for (int i = 0; i < renderLineArray.length(); i++)
    {
        segmentArray.append(renderLineArray.renderLine(i, status).getLine().length());
    }
    return segmentArray;
}

MObjectArray PfxHairHelper::GetHairBaseMesh()
{
    return GetTypeNodes(m_hairSystem, MItDependencyGraph::kUpstream, MFn::kMesh);
}

IndexArrayMap PfxHairHelper::GetMeshHairIndexMap()
{
    MStatus status;
    IndexArrayMap meshHairIndexMap;
    MFnDependencyNode fnHairSystem(m_hairSystem);
    MPlug inputHairPlug = fnHairSystem.findPlug("inputHair");
    // AiMsgInfo("[%s] : inputHair array count : %d", kTranslatorName, inputHairPlug.numElements());
    for (unsigned int i = 0; i < inputHairPlug.numElements(); i++)
    {
        MPlug inputHairIndexPlug = inputHairPlug.connectionByPhysicalIndex(i, &status);
        MPlug folliclePlug = inputHairIndexPlug.source(&status);
        if (folliclePlug.isNull() || status == MS::kFailure)
            continue;
        MFnDependencyNode fnFollicleNode(folliclePlug.node());
        MPlug inMeshPlug = fnFollicleNode.findPlug("inputMesh");
        MPlug meshPlug = inMeshPlug.source(&status);
        // AiMsgInfo("[%s] : inputMesh plug : %s", kTranslatorName , meshPlug.name().asChar());
        if (meshPlug.isNull() || status == MS::kFailure)
            continue;
        MObject meshObject = meshPlug.node();
        MFnMesh fnMesh(meshObject);
        string meshPath = fnMesh.fullPathName().asChar();
        // AiMsgInfo("[%s] : inputMesh : %s", kTranslatorName, meshPath.c_str());
        if (meshHairIndexMap.find(meshPath) == meshHairIndexMap.end())
            meshHairIndexMap[meshPath] = MIntArray();
        meshHairIndexMap[meshPath].append(i);
    }
    return meshHairIndexMap;
}

MStatus PfxHairHelper::GetMeshFromHairIndex(int hairIndex, MObject &hairMesh)
{
    IndexArrayMap meshHairIndexMap = GetMeshHairIndexMap();
    if (meshHairIndexMap.empty())
        return MS::kNotFound;
    // for (const auto &[mesh, indexes] : meshHairIndexMap)
    // for (auto indexMap : meshHairIndexMap)
    for (auto it = meshHairIndexMap.begin(); it != meshHairIndexMap.end(); ++it)
    {
        string mesh = it->first;
        MIntArray indexes = it->second;
        for (unsigned int i = 0; i < indexes.length(); i++)
        {
            if (hairIndex != indexes[i])
                continue;
            MSelectionList s_list;
            s_list.clear();
            s_list.add(mesh.c_str());
            MObject meshObject;
            s_list.getDependNode(0, meshObject);
            hairMesh = meshObject;
            return MS::kSuccess;
        }
    }
    return MS::kNotFound;
}

MStatus PfxHairHelper::GetMeshFromHairIndex(int hairIndex, MString &hairMesh)
{
    MObject meshObject;
    GetMeshFromHairIndex(hairIndex, meshObject);
    if (meshObject.isNull())
        return MS::kNotFound;
    hairMesh = MFnDagNode(meshObject).fullPathName();
    return MS::kSuccess;
}

void PfxHairHelper::GetHairRootPoints(PointArrayMap &rootPointArrayMap)
{
    MStatus status;

    MRenderLineArray renderLineArray = GetHairData();
    int renderLineCount = renderLineArray.length();
    int hairPerClump = GetHairsPerClump();
    AiMsgInfo("[%s] : Hair Render Line Count : %d", kTranslatorName, renderLineCount);
    AiMsgInfo("[%s] : Hair Per Clump : %d", kTranslatorName, hairPerClump);
    for (auto i = 0; i < renderLineArray.length(); i++)
    {
        int hairIndex = i / hairPerClump;
        MRenderLine renderLine = renderLineArray.renderLine(i, &status);
        MPoint rootPoint(renderLine.getLine()[0]);
        MString hairMesh;
        if (GetMeshFromHairIndex(hairIndex, hairMesh) == MS::kNotFound)
            continue;
        AiMsgInfo("[%s] : Hair Mesh Index : [%d] : %s", kTranslatorName, hairIndex, hairMesh.asChar());
        rootPointArrayMap[hairMesh.asChar()].append(rootPoint);
    }
}

void PfxHairHelper::GetHairRootPointsAndNormals(MPointArray &pointArray, MVectorArray &vectorArray, bool debug)
{
    MStatus status;
    PointArrayMap rootPointArrayMap;
    GetHairRootPoints(rootPointArrayMap);
    // for (auto rp : rootPointArrayMap)
    for (auto it = rootPointArrayMap.begin(); it != rootPointArrayMap.end(); ++it)
    {
        string mesh = it->first;
        MPointArray rootPoints = it->second;

        MDagPath dagPath;
        GetMayaDagPath(mesh, dagPath);

        if (!dagPath.isValid() || !dagPath.hasFn(MFn::kMesh))
        {
            MGlobal::displayError(MString(mesh.c_str()) + " is invalid");
            continue;
        }

        MMeshIntersector intersector;
        MMatrix matrix = dagPath.inclusiveMatrix();
        status = intersector.create(dagPath.node(), matrix);
        for (unsigned i = 0; i < rootPoints.length(); i++)
        {
            MPoint pt = rootPoints[i];
            MPointOnMesh pointInfo;
            intersector.getClosestPoint(pt, pointInfo);
            MFloatVector closestNormal = pointInfo.getNormal();
            closestNormal = closestNormal * -1;
            pointArray.append(pt);
            vectorArray.append(closestNormal);
        }
    }
}

int PfxHairHelper::GetHairIndexFromCurveId(int curveId, MStatus *status = nullptr)
{
    int hairPerClump = GetHairsPerClump();
    int hairCount = GetHairCount();
    if (curveId > hairCount)
    {
        // std::cerr << "[Error] : curveId (" << curveId << ") is large the hair count( "
        //           << hairCount << ")" << endl;
        return 0;
    }
    return curveId / hairPerClump;
}

bool PfxHairHelper::IsRenderable()
{
    MFnDagNode fnPfxHair(m_pfxHair);
    MDagPath pfxHairPath;
    GetMayaDagPath(fnPfxHair.fullPathName().asChar(), pfxHairPath);

    if (pfxHairPath.isVisible())
    {
        MFnDagNode fnHairSystem(m_hairSystem);
        MPlug shaderPlug = fnHairSystem.findPlug("aiHairShader");
        MObjectArray shaderNodes = GetTypeNodes(shaderPlug, MItDependencyGraph::kUpstream);
        MString sampleShaderType("apCurveSample");
        for (unsigned int i = 0; i < shaderNodes.length(); i++)
        {
            MFnDependencyNode fnNode(shaderNodes[i]);
            AiMsgDebug("[%s] : shader : %s", kTranslatorName, fnNode.typeName().asChar());
            if (fnNode.typeName() == sampleShaderType)
            {
                AiMsgDebug("[%s] : Found %s : %s", kTranslatorName, sampleShaderType.asChar(), fnNode.typeName().asChar());
                return true;
            }
        }
    }
    return false;
}

int PfxHairHelper::GetHairsPerClump()
{
    MFnDependencyNode fnHairNode(m_hairSystem);
    MPlug hpcPlug = fnHairNode.findPlug("hairsPerClump");
    return hpcPlug.asInt();
}