# Copyright (c) 2017-2026 Soft8Soft, LLC. All rights reserved.
#
# Use of this software is subject to the terms of the Verge3D license
# agreement provided at the time of installation or download, or which
# otherwise accompanies this software in either electronic or hard copy form.
import pluginUtils
log = pluginUtils.log.getLogger('V3D-MY')

import maya.cmds as cmds
import maya.OpenMaya as OpenMaya

import customAttrs
import mayaUtils
import matNodes

class Collector():

    defaultCameras = ['|top', '|front', '|side', '|persp']

    def __init__(self, exportSettings):

        self.exportSettings = exportSettings;

        self.nodes = []
        self.locatorNodes = []
        self.topLevelNodes = []
        self.meshes = []
        self.curves = [] # curves here are special kind of meshes
        self.cameras = []
        self.lights = []
        self.lightProbes = []
        self.clippingPlanes = []
        self.skyDomeLights = []
        self.materials = []
        self.textures = []

        self.materialTextures = []
        self.convertedMeshes = []

    def collect(self):

        nodes = cmds.ls(selection=self.exportSettings['selectionOnly'], assemblies=True, long=True)

        for node in nodes:
            self.collectNode(node, True)

        for node in (self.meshes + self.curves):
            if not (mayaUtils.isNurbsCurve(node) or mayaUtils.isBezierCurve(node)):
                shadingGrps = cmds.listConnections(node, type='shadingEngine')
                for mat in cmds.ls(cmds.listConnections(shadingGrps), materials=True):
                    self.collectMat(mat, node)

        for mat in self.materials:

            self.materialTextures.append([])

            matSettings = customAttrs.parseMaterialSettings(mat)

            if matSettings['gltfCompat']:
                matType = cmds.objectType(mat)

                if matType == 'aiStandardSurface' or matType == 'standardSurface':
                    tex = mayaUtils.extractTexture(mat, 'baseColor')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

                    metallicTex = mayaUtils.extractTexture(mat, 'metalness')
                    roughnessTex = mayaUtils.extractTexture(mat, 'specularRoughness')
                    if metallicTex and metallicTex == roughnessTex:
                        self.textures.append(metallicTex)
                        self.materialTextures[-1].append(metallicTex)

                    occlusionTex1 = mayaUtils.extractTexture(mat, 'base')
                    occlusionTex2 = mayaUtils.extractTexture(mat, 'specular')
                    if occlusionTex1 and occlusionTex1 == occlusionTex2:
                        self.textures.append(occlusionTex1)
                        self.materialTextures[-1].append(occlusionTex1)

                    tex = mayaUtils.extractTexture(mat, 'normalCamera', 'bumpValue')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

                    tex = mayaUtils.extractTexture(mat, 'emissionColor')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

                elif matType == 'lambert':
                    tex = mayaUtils.extractTexture(mat, 'color')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

                elif matType == 'usdPreviewSurface':
                    tex = mayaUtils.extractTexture(mat, 'diffuseColor')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

                    metallicTex = mayaUtils.extractTexture(mat, 'metallic')
                    roughnessTex = mayaUtils.extractTexture(mat, 'roughness')
                    if metallicTex and metallicTex == roughnessTex:
                        self.textures.append(metallicTex)
                        self.materialTextures[-1].append(metallicTex)

                    tex = mayaUtils.extractTexture(mat, 'emissiveColor')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

                    tex = mayaUtils.extractTexture(mat, 'occlusion')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

                    tex = mayaUtils.extractTexture(mat, 'normal')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

                elif matType == 'StingrayPBS':
                    tex = mayaUtils.extractTexture(mat, 'TEX_color_map')
                    if tex:
                        self.textures.append(tex)
                        self.materialTextures[-1].append(tex)

            else:
                nodeTextures = []
                matNodes.extractTextures(matNodes.nodeGraphSG(mat), nodeTextures)
                self.textures.extend(nodeTextures)
                self.materialTextures[-1].extend(nodeTextures)

        if len(self.skyDomeLights):
            nodeTextures = []
            matNodes.extractTextures(self.skyDomeLights[0], nodeTextures)
            self.textures.extend(nodeTextures)
            self.materialTextures.append(nodeTextures)

        if len(self.cameras) == 0:
            node = '|persp'

            childNodes = cmds.listRelatives(node, children=True, fullPath=True)
            if childNodes:
                for child in childNodes:
                    if mayaUtils.isCamera(child):
                        self.cameras.append(child)
                        self.nodes.append(node)
                        self.topLevelNodes.append(node)
        else:
            activeCamera = None

            for vp in cmds.getPanel(type='modelPanel'):
                activeCamera = cmds.modelEditor(vp, query=True, activeView=True, camera=True)

            if activeCamera:
                for cam in self.cameras:
                    if mayaUtils.getFirstShape(activeCamera) == cam:
                        self.cameras.remove(cam)
                        self.cameras.insert(0, cam)
                        break

        shadowMeshes = []
        for m in self.meshes + self.curves:
            if cmds.getAttr(m + '.receiveShadows') or cmds.getAttr(m + '.castsShadows'):
                shadowMeshes.append(m)

        if shadowMeshes:
            self.shadowBox = cmds.exactWorldBoundingBox(shadowMeshes)
        else:
            self.shadowBox = [-100, -100, -100, 100, 100, 100]

    def collectNode(self, node, topLevel):
        if node in self.defaultCameras:
            return

        node = self.convertNode(node)

        self.nodes.append(node)

        if topLevel:
            self.topLevelNodes.append(node)

        childNodes = cmds.listRelatives(node, children=True, fullPath=True)
        if childNodes:
            for child in childNodes:
                if not self.exportSettings['bakeText'] and self.exportSettings['format'] != 'HTML' and mayaUtils.isMeshType(child):
                    self.curves.append(child)
                elif (mayaUtils.isMesh(child) or
                        mayaUtils.isNurbsCurve(child) or
                        mayaUtils.isBezierCurve(child)):
                    self.meshes.append(child)
                elif mayaUtils.isCamera(child):
                    self.cameras.append(child)
                elif mayaUtils.isLight(child):
                    self.lights.append(child)
                elif mayaUtils.isLightProbe(child):
                    self.lightProbes.append(child)
                elif mayaUtils.isClippingPlane(child):
                    self.clippingPlanes.append(child)
                elif mayaUtils.isSkyDomeLight(child):
                    self.skyDomeLights.append(child)
                elif mayaUtils.isLocatorNode(child):
                    self.locatorNodes.append(child)
                elif mayaUtils.isTransformNode(child):
                    self.collectNode(child, False)
                elif mayaUtils.isJoint(child):
                    self.collectNode(child, False)

    def convertNode(self, node):
        """Returns new node if conversion needed"""

        childNodes = cmds.listRelatives(node, children=True, fullPath=True)
        if childNodes:
            shape = childNodes[0]

            if mayaUtils.isNurbsSurface(shape):

                surfName = mayaUtils.getName(node)
                log.info('Converting {} surface to poly mesh'.format(surfName))
                nodeConv = cmds.nurbsToPoly(node,
                                            name = surfName + 'Converted',
                                            constructionHistory = True,
                                            format = 1,
                                            polygonType = 1,
                                            fractionalTolerance = 0.01,
                                            minEdgeLength = 0.001,
                                            delta = 0.1,
                                            useSurfaceShader = True)[0]

                nodeParent = cmds.listRelatives(node, parent=True, fullPath=True)
                if nodeParent:
                    nodeConv = cmds.ls(cmds.parent(nodeConv, nodeParent[0]), long=True)[0]

                sets = mayaUtils.extractObjectSetNames(node)
                for s in sets:
                    cmds.sets(nodeConv, include=s)

                self.convertedMeshes.append(nodeConv)

                node = nodeConv

        return node

    def collectMat(self, mat, node):
        if mat and mat not in self.materials:
            self.materials.append(mat)

    def cleanup(self):
        for node in self.convertedMeshes:
            cmds.delete(node)
