#include "CurveDataTranslator.h"
#include "Xgen.h"
#include "PfxHair.h"

#include <time.h>
#include <ai.h>
#include <ai_msg.h>
#include <ai_nodes.h>

#include <maya/MFnSet.h>
#include <maya/MFnAttribute.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MFnPfxGeometry.h>
#include <maya/MRenderLine.h>
#include <maya/MRenderLineArray.h>
#include <maya/MDagPath.h>
#include <maya/MPxData.h>
#include <maya/MSelectionList.h>
#include <maya/MFnPluginData.h>
#include <maya/MItDependencyGraph.h>
#include <maya/MItDependencyNodes.h>
#include <maya/MObject.h>
#include <maya/MObjectArray.h>
#include <maya/MPlug.h>

void CCurveDataTranslator::NodeInitializer(CAbTranslator context)
{
    AiMsgInfo("[%s] : NodeInitializer()", kTranslatorName);
}

AtNode *CCurveDataTranslator::CreateArnoldNodes()
{
    AiMsgInfo("[%s] : CreateArnoldNodes()", kTranslatorName);
    return AddArnoldNode("apCurveSample");
}

void CCurveDataTranslator::Delete()
{
}

void CCurveDataTranslator::Update(AtNode *shader)
{
}

void CCurveDataTranslator::Export(AtNode *shader)
{
    clock_t tic, toc;
    tic = clock();
    MStatus status;

    AtNode *option = AiUniverseGetOptions();

    AiMsgInfo("[%s] : Export()", kTranslatorName);
    AiMsgInfo("[%s] : Starting translate %s", kTranslatorName, GetMayaNodeName().asChar());
    MFnDependencyNode fnDepNode(GetMayaObject());
    MPlug outColorPlug = fnDepNode.findPlug("outColor", true, &status);

    MPlugArray destPlugArray;
    outColorPlug.destinations(destPlugArray, &status);

    MItDependencyNodes dgHairIt(MFn::kPfxHair);

    AiMsgInfo("[%s] : Looking pfxHair ...", kTranslatorName);
    for (; !dgHairIt.isDone(); dgHairIt.next())
    {

        MObject currentNode = dgHairIt.thisNode();

        PfxHairHelper pfxHairHelper(currentNode);
        MFnDagNode fnPfxHairNode(currentNode);
        AiMsgInfo("[%s] : Current Node : %s", kTranslatorName, fnPfxHairNode.fullPathName().asChar());

        if (!pfxHairHelper.IsRenderable())
            continue;

        MIntArray hairSegments = pfxHairHelper.GetHairBaseSegement();

        int cvCount = hairSegments[0];
        MString pfxHairPath = fnPfxHairNode.fullPathName();
        pfxHairPath.substitute("|", "/");

        const MString cvCountParam = MString(pfxHairPath + "_cv_count");

        if (!AiNodeLookUpUserParameter(option, cvCountParam.asChar()))
            AiNodeDeclare(option, cvCountParam.asChar(), "constant INT");

        AiNodeSetInt(option, AtString(cvCountParam.asChar()), cvCount);
        AiMsgInfo("[%s] : Export params : %s", kTranslatorName, cvCountParam.asChar());

        AiMsgInfo("-------------------------------------");
        AiMsgInfo("[%s] : pfxHair : %s", kTranslatorName, fnPfxHairNode.partialPathName().asChar());
        AiMsgInfo("-------------------------------------");

        MObjectArray hairBaseMeshList = pfxHairHelper.GetHairBaseMesh();
        const MString meshParam = MString(pfxHairPath + "_meshs");

        if (!AiNodeLookUpUserParameter(option, meshParam.asChar()))
            AiNodeDeclare(option, meshParam.asChar(), "constant ARRAY STRING");

        AtArray *meshArray = AiArrayAllocate(hairBaseMeshList.length(), 1, AI_TYPE_STRING);
        for (unsigned int i = 0; i < hairBaseMeshList.length(); i++)
        {
            MFnDependencyNode meshNode(hairBaseMeshList[i]);
            AiArraySetStr(meshArray, i, AtString(meshNode.name().asChar()));
        }
        AiNodeSetArray(option, meshParam.asChar(), meshArray);

        AiMsgInfo("[%s] : array string[%d] : %s", kTranslatorName, hairBaseMeshList.length(), meshParam.asChar());

        MPointArray rootPoints;
        MVectorArray rootNormals;
        pfxHairHelper.GetHairRootPointsAndNormals(rootPoints, rootNormals);

        AiMsgInfo("[%s] : Export root points", kTranslatorName);
        const MString rootPointParam = MString(pfxHairPath + "_points");

        AiMsgInfo("[%s] : array vector[%d] : %s", kTranslatorName, rootPoints.length(), rootPointParam.asChar());

        if (!AiNodeLookUpUserParameter(option, rootPointParam.asChar()))
            AiNodeDeclare(option, rootPointParam.asChar(), "constant ARRAY VECTOR");

        AtArray *pointArray = AiArrayAllocate(rootPoints.length(), 1, AI_TYPE_VECTOR);

        for (unsigned i = 0; i < rootPoints.length(); i++)
        {
            MPoint rp = rootPoints[i];
            AiArraySetVec(pointArray, i, AtVector(rp.x, rp.y, rp.z));
        }

        AiNodeSetArray(option, rootPointParam.asChar(), pointArray);
        AiMsgInfo("-------------------------------------");

        AiMsgInfo("[%s] : Export root normals", kTranslatorName);

        const MString rootNormalParam = MString(pfxHairPath + "_normals");

        AiMsgInfo("[%s] : array vector[%d] : %s", kTranslatorName, rootNormals.length(), rootNormalParam.asChar());

        if (!AiNodeLookUpUserParameter(option, rootNormalParam.asChar()))
            AiNodeDeclare(option, rootNormalParam.asChar(), "constant ARRAY VECTOR");
        AtArray *normalArray = AiArrayAllocate(rootNormals.length(), 1, AI_TYPE_VECTOR);

        for (unsigned i = 0; i < rootNormals.length(); i++)
        {
            MVector rn = rootNormals[i];
            AiArraySetVec(normalArray, i, AtVector(rn.x, rn.y, rn.z));
        }

        AiNodeSetArray(option, rootNormalParam.asChar(), normalArray);
        AiMsgInfo("-------------------------------------");
    }

    // Looking the shader assignment
    MItDependencyGraph *dgIt;
    dgIt = new MItDependencyGraph(outColorPlug, MFn::kShadingEngine,
                                  MItDependencyGraph::kDownstream,
                                  MItDependencyGraph::kBreadthFirst,
                                  MItDependencyGraph::kNodeLevel,
                                  &status);
    for (; !dgIt->isDone(); dgIt->next())
    {
        // Get shading group from current item
        MObject dest = dgIt->currentItem();

        if (!dest.hasFn(MFn::kShadingEngine))
            continue;

        MFnSet fnSG(dest);
        MSelectionList setMember;
        fnSG.getMembers(setMember, true);
        // Look each assignment of the SG
        for (unsigned j = 0; j < setMember.length(); j++)
        {
            MObject desc;
            MDagPath descPath;

            setMember.getDependNode(j, desc);
            setMember.getDagPath(j, descPath);
            MFnDagNode fnDescNode(descPath);

            XGenHelper xgenHelper(desc);

            if (!descPath.isVisible())
                continue;
            MString descPathName = fnDescNode.fullPathName();
            descPathName.substitute("|", "/");
            MString descParamBaseName = descPathName + "_curves";
            AiMsgInfo("[%s] : %s api_type=%s", kTranslatorName, fnDescNode.name().asChar(), desc.apiTypeStr());
            AiMsgInfo("[%s] : %s node_type=%s", kTranslatorName, fnDescNode.name().asChar(), fnDescNode.typeName().asChar());
            // XGen Interactive Groom Spline
            // We only care about xgmSplineDescription
            if (fnDescNode.typeName() == "xgmSplineDescription")
            {
                // Set Export RenderStats
                fnDescNode.findPlug("castsShadows").setValue(0);
                fnDescNode.findPlug("receiveShadows").setValue(0);
                fnDescNode.findPlug("aiVisibleInDiffuseReflection").setValue(0);
                fnDescNode.findPlug("aiVisibleInSpecularReflection").setValue(0);

                static const MString sApplyRenderOverrideCmd = "xgmSplineApplyRenderOverride ";
                MGlobal::executeCommand(sApplyRenderOverrideCmd + fnDescNode.partialPathName());

                MObject splineBase = xgenHelper.GetSplineBase();
                // GetSplineBaseNode(descObject, splineBase);

                MFnDependencyNode fnSplineBase(splineBase);
                // int cvCount = GetSplineBaseSegement(splineBase);
                int cvCount = xgenHelper.GetBaseSegment();

                const MString cvCountParam = descParamBaseName + "_cv_count";
                if (!AiNodeLookUpUserParameter(option, cvCountParam.asChar()))
                    AiNodeDeclare(option, cvCountParam.asChar(), "constant INT");
                AiNodeSetInt(option, AtString(cvCountParam.asChar()), cvCount);

                // MSelectionList meshList = GetSplineBaseMesh(splineBase);
                MObjectArray meshList = xgenHelper.GetSplineBaseMesh();

                AiMsgInfo("-------------------------------------");
                AiMsgInfo("[%s] : xgmSplineDescription : %s", kTranslatorName, fnDescNode.partialPathName().asChar());
                AiMsgInfo("[%s] : xgmSplineBase : %s", kTranslatorName, fnSplineBase.name().asChar());
                AiMsgInfo("-------------------------------------");

                const MString meshParam = descParamBaseName + "_meshs";

                if (!AiNodeLookUpUserParameter(option, meshParam.asChar()))
                    AiNodeDeclare(option, meshParam.asChar(), "constant ARRAY STRING");

                AtArray *meshArray = AiArrayAllocate(meshList.length(), 1, AI_TYPE_STRING);
                for (unsigned int i = 0; i < meshList.length(); i++)
                {
                    MObject meshObj = meshList[i];
                    MFnDagNode meshNode(meshObj);
                    MString meshPath = meshNode.fullPathName();
                    meshPath.substitute("|", "/");
                    AiArraySetStr(meshArray, i, AtString(meshPath.asChar()));
                }

                AiNodeSetArray(option, meshParam.asChar(), meshArray);

                MPointArray rootPoints;
                MVectorArray rootNormals;

                // GetSplineRootPointsAndNormals(fnDescNode.partialPathName().asChar(), rootPoints, rootNormals);
                xgenHelper.GetSplineRootPointsAndNormals(rootPoints, rootNormals);

                // getSplineRootPoints(fnDescNode.partialPathName().asChar(), rootPoints);

                AiMsgInfo("[%s] : Export root points", kTranslatorName);
                const MString rootPointParam = descParamBaseName + "_points";

                AiMsgInfo("[%s] : array vector[%d] : %s", kTranslatorName, rootPoints.length(), rootPointParam.asChar());

                if (!AiNodeLookUpUserParameter(option, rootPointParam.asChar()))
                    AiNodeDeclare(option, rootPointParam.asChar(), "constant ARRAY VECTOR");

                AtArray *pointArray = AiArrayAllocate(rootPoints.length(), 1, AI_TYPE_VECTOR);

                for (unsigned i = 0; i < rootPoints.length(); i++)
                {
                    MPoint rp = rootPoints[i];
                    AiArraySetVec(pointArray, i, AtVector(rp.x, rp.y, rp.z));
                }

                AiNodeSetArray(option, rootPointParam.asChar(), pointArray);

                AiMsgInfo("-------------------------------------");

                // getSplineRootNormals(fnDescNode.partialPathName().asChar(), rootNormals);
                AiMsgInfo("[%s] : Export root normals", kTranslatorName);

                const MString rootNormalParam = descParamBaseName + "_normals";

                AiMsgInfo("[%s] : array vector[%d] : %s", kTranslatorName, rootNormals.length(), rootNormalParam.asChar());

                if (!AiNodeLookUpUserParameter(option, rootNormalParam.asChar()))
                    AiNodeDeclare(option, rootNormalParam.asChar(), "constant ARRAY VECTOR");
                AtArray *normalArray = AiArrayAllocate(rootNormals.length(), 1, AI_TYPE_VECTOR);

                for (unsigned i = 0; i < rootNormals.length(); i++)
                {
                    MVector rn = rootNormals[i];
                    AiArraySetVec(normalArray, i, AtVector(rn.x, rn.y, rn.z));
                }

                AiNodeSetArray(option, rootNormalParam.asChar(), normalArray);
                AiMsgInfo("-------------------------------------");
            }
        }
    }

    AiMsgInfo("[%s] : Exporting %s settings", kTranslatorName, GetMayaNodeName().asChar());
    ProcessParameter(shader, "sample_data_type", AI_TYPE_INT, "sampleDataType");
    ProcessParameter(shader, "normalize", AI_TYPE_BOOLEAN, "normalize");
    ProcessParameter(shader, "trace_set", AI_TYPE_STRING, "traceSet");
    ProcessParameter(shader, "ray_offset", AI_TYPE_FLOAT, "rayOffset");
    ProcessParameter(shader, "ray_max_depth", AI_TYPE_INT, "rayMaxDepth");
    ProcessParameter(shader, "ray_max_sample", AI_TYPE_INT, "rayMaxSample");
    ProcessParameter(shader, "occ_min_dist", AI_TYPE_FLOAT, "occMinDist");
    ProcessParameter(shader, "occ_max_dist", AI_TYPE_FLOAT, "occMaxDist");
    ProcessParameter(shader, "occ_spread", AI_TYPE_FLOAT, "occSpread");
    ProcessParameter(shader, "occ_falloff", AI_TYPE_FLOAT, "occFalloff");
    ProcessParameter(shader, "error_color", AI_TYPE_RGB, "errorColor");
    ProcessParameter(shader, "verbose", AI_TYPE_BOOLEAN, "verbose");

    // Export parameter to shader tree or value
    /*
    ProcessParameter(shader, "root_color_gain", AI_TYPE_RGB, "rootColorGain");
    ProcessParameter(shader, "root_color_offset", AI_TYPE_RGB, "rootColorOffset");
    ProcessParameter(shader, "tip_color_gain", AI_TYPE_RGB, "tipColorGain");
    ProcessParameter(shader, "tip_color_offset", AI_TYPE_RGB, "tipColorOffset");
    ProcessParameter(shader, "ray_offset", AI_TYPE_FLOAT, "rayOffset");
    ProcessParameter(shader, "ray_max_depth", AI_TYPE_INT, "rayMaxDepth");
    ProcessParameter(shader, "ray_max_sample", AI_TYPE_INT, "rayMaxSample");
    ProcessParameter(shader, "trace_set", AI_TYPE_STRING, "traceSet");
    ProcessParameter(shader, "verbose", AI_TYPE_BOOLEAN, "verbose");
    ProcessParameter(shader, "debug", AI_TYPE_BOOLEAN, "debug");
    */

    toc = clock();

    AiMsgInfo("[%s] : Export in %lf sec", kTranslatorName, (toc - tic) / (double)CLOCKS_PER_SEC);
    AiMsgInfo("[%s] : Export finish : %s", kTranslatorName, fnDepNode.name().asChar());
}