"""Shifters Rig Main class."""
import os.path
import datetime
import getpass

# Maya
import pymel.core as pm
from pymel.core import datatypes
from pymel import versions

# mgear
import mgear
import mgear.maya.utils
from . import guide, component

from mgear.maya import primitive, attribute, skin, dag, icon


# check if we have loaded the necessary plugins
if not pm.pluginInfo("mgear_solvers", q=True, loaded=True):
    try:
        pm.loadPlugin("mgear_solvers")
    except RuntimeError:
        pm.displayError("You need the mgear_solvers plugin!")
if not pm.pluginInfo("matrixNodes", q=True, loaded=True):
    pm.loadPlugin("matrixNodes")

COMPONENT_PATH = os.path.join(os.path.dirname(__file__), "component")
TEMPLATE_PATH = os.path.join(COMPONENT_PATH, "templates")
SYNOPTIC_PATH = os.path.abspath(os.path.join(
    os.path.dirname(__file__), os.pardir, "synoptic", "tabs"))

SHIFTER_COMPONENT_ENV_KEY = "MGEAR_SHIFTER_COMPONENT_PATH"


def getComponentDirectories():
    """Get the components directory"""
    return mgear.maya.utils.gatherCustomModuleDirectories(
        SHIFTER_COMPONENT_ENV_KEY,
        os.path.join(os.path.dirname(__file__), "component"))


def importComponentGuide(comp_type):
    """Import the Component guide"""
    dirs = getComponentDirectories()
    defFmt = "mgear.maya.shifter.component.{}.guide"
    customFmt = "{}.guide"

    module = mgear.maya.utils.importFromStandardOrCustomDirectories(
        dirs, defFmt, customFmt, comp_type)
    return module


def importComponent(comp_type):
    """Import the Component """
    dirs = getComponentDirectories()
    defFmt = "mgear.maya.shifter.component.{}"
    customFmt = "{}"

    module = mgear.maya.utils.importFromStandardOrCustomDirectories(
        dirs, defFmt, customFmt, comp_type)
    return module


def reloadComponents(*args):
    """Reload all componets

    Args:
        *args: Dummy
    """
    compDir = getComponentDirectories()

    for x in compDir:
        for com in compDir[x]:
            try:
                reload(importComponent(com))
                reload(importComponentGuide(com))
                print "reload : {}.{}".format(os.path.basename(x), com)
            except ImportError:
                pass


class Rig(object):
    """The main rig class.

    Attributes:
        guide: guide.Rig() initialization.
        groups (dic): Rig groups (Maya sets)
        components (dic): Dictionary for the rig components.
            Keys are the component fullname (ie. 'arm_L0')
        componentsIndex (list): Components index list.

    """

    def __init__(self):

        self.guide = guide.Rig()

        self.groups = {}
        self.subGroups = {}

        self.components = {}
        self.componentsIndex = []

        self.customStepDic = {}

    def buildFromSelection(self):
        """Build the rig from selected guides."""

        startTime = datetime.datetime.now()
        mgear.log("\n" + "= SHIFTER RIG SYSTEM " + "=" * 46)

        self.stopBuild = False
        selection = pm.ls(selection=True)
        if not selection:
            mgear.log(
                "Select one or more guide root or a guide model",
                mgear.sev_error)
            return

        self.preCustomStep(selection)

        if not self.stopBuild:
            mgear.log("\n" + "= GUIDE VALIDATION " + "=" * 46)
            # Check guide is valid
            self.guide.setFromSelection()
            if not self.guide.valid:
                return

            # Build
            mgear.log("\n" + "= BUILDING RIG " + "=" * 46)
            self.build()
            self.postCustomStep()

            endTime = datetime.datetime.now()
            finalTime = endTime - startTime
            mgear.log("\n" + "= SHIFTER BUILD RIG DONE {} [ {} ] {}".format(
                "=" * 16,
                finalTime,
                "=" * 7
            ))

    def build(self):
        """Build the rig."""

        self.options = self.guide.values
        self.guides = self.guide.components

        self.customStepDic["mgearRun"] = self

        self.initialHierarchy()
        self.processComponents()
        self.finalize()

        return self.model

    def stepsList(self, checker, attr):
        if self.options[checker] and self.options[attr]:
            return self.options[attr].split(",")
        else:
            return None

    def customStep(self, customSteps=None):
        if customSteps:
            for step in customSteps:
                if not self.stopBuild:
                    self.stopBuild = guide.helperSlots.runStep(
                        step.split("|")[-1][1:], self.customStepDic)
                else:
                    pm.displayWarning("Build Stopped")
                    break

    def preCustomStep(self, selection):
        if (selection[0].hasAttr("ismodel") and
                selection[0].attr("doPreCustomStep").get()):
            customSteps = selection[0].attr("preCustomStep").get()
            if customSteps:
                mgear.log("\n" + "= PRE CUSTOM STEPS " + "=" * 46)
                self.customStep(customSteps.split(","))

    def postCustomStep(self):
        customSteps = self.stepsList("doPostCustomStep", "postCustomStep")
        if customSteps:
            mgear.log("\n" + "= POST CUSTOM STEPS " + "=" * 46)
            self.customStep(customSteps)

    def initialHierarchy(self):
        """Build the initial hierarchy of the rig.

        Create the rig model, the main properties,
        and a couple of base organisation nulls.
        Get the global size of the rig.

        """
        mgear.log("Initial Hierarchy")

        # --------------------------------------------------
        # Model
        self.model = primitive.addTransformFromPos(
            None, self.options["rig_name"])
        attribute.lockAttribute(self.model)

        # ------------------------- -------------------------
        # Global Ctl
        if self.options["worldCtl"]:
            self.global_ctl = self.addCtl(self.model,
                                          "world_ctl",
                                          datatypes.Matrix(),
                                          self.options["C_color_fk"],
                                          "circle", w=10)
        else:
            self.global_ctl = self.addCtl(self.model,
                                          "global_C0_ctl",
                                          datatypes.Matrix(),
                                          self.options["C_color_fk"],
                                          "crossarrow",
                                          w=10)
        attribute.setRotOrder(self.global_ctl, "ZXY")

        # --------------------------------------------------
        # Setup in world Space
        self.setupWS = primitive.addTransformFromPos(self.model, "setup")
        attribute.lockAttribute(self.setupWS)

        # --------------------------------------------------
        # INFOS
        self.isRig_att = attribute.addAttribute(
            self.model, "is_rig", "bool", True)
        self.rigName_att = attribute.addAttribute(
            self.model, "rig_name", "string", self.options["rig_name"])
        self.user_att = attribute.addAttribute(
            self.model, "user", "string", getpass.getuser())
        self.isWip_att = attribute.addAttribute(
            self.model, "wip", "bool", self.options["mode"] != 0)
        self.date_att = attribute.addAttribute(
            self.model, "date", "string", str(datetime.datetime.now()))
        self.mayaVersion_att = attribute.addAttribute(
            self.model, "maya_version", "string",
            str(pm.mel.eval("getApplicationVersionAsFloat")))
        self.gearVersion_att = attribute.addAttribute(
            self.model, "gear_version", "string", mgear.getVersion())
        self.synoptic_att = attribute.addAttribute(
            self.model, "synoptic", "string", str(self.options["synoptic"]))
        self.comments_att = attribute.addAttribute(
            self.model, "comments", "string", str(self.options["comments"]))
        self.ctlVis_att = attribute.addAttribute(
            self.model, "ctl_vis", "bool", True)
        self.ctlVisPlayback_att = attribute.addAttribute(
            self.model, "ctl_vis_on_playback", "bool", True)
        self.jntVis_att = attribute.addAttribute(
            self.model, "jnt_vis", "bool", True)

        self.qsA_att = attribute.addAttribute(
            self.model, "quickselA", "string", "")
        self.qsB_att = attribute.addAttribute(
            self.model, "quickselB", "string", "")
        self.qsC_att = attribute.addAttribute(
            self.model, "quickselC", "string", "")
        self.qsD_att = attribute.addAttribute(
            self.model, "quickselD", "string", "")
        self.qsE_att = attribute.addAttribute(
            self.model, "quickselE", "string", "")
        self.qsF_att = attribute.addAttribute(
            self.model, "quickselF", "string", "")

        self.rigGroups = self.model.addAttr("rigGroups", at='message', m=1)
        self.rigPoses = self.model.addAttr("rigPoses", at='message', m=1)

        # Connect global visibility
        pm.connectAttr(self.ctlVis_att, self.global_ctl.attr("visibility"))
        pm.connectAttr(self.ctlVisPlayback_att,
                       self.global_ctl.attr("hideOnPlayback"))
        attribute.lockAttribute(self.global_ctl, ['v'])

        # --------------------------------------------------
        # Basic set of null
        if self.options["joint_rig"]:
            self.jnt_org = primitive.addTransformFromPos(self.model, "jnt_org")
            pm.connectAttr(self.jntVis_att, self.jnt_org.attr("visibility"))

    def processComponents(self):
        """
        Process the components of the rig, following the creation steps.
        """

        # Init
        self.components_infos = {}

        for comp in self.guide.componentsIndex:
            guide_ = self.guides[comp]
            mgear.log("Init : " + guide_.fullName + " (" + guide_.type + ")")

            module = importComponent(guide_.type)
            Component = getattr(module, "Component")

            comp = Component(self, guide_)
            if comp.fullName not in self.componentsIndex:
                self.components[comp.fullName] = comp
                self.componentsIndex.append(comp.fullName)

                self.components_infos[comp.fullName] = [
                    guide_.compType, guide_.getVersion(), guide_.author]

        # Creation steps
        self.steps = component.Main.steps
        for i, name in enumerate(self.steps):
            # for count, compName in enumerate(self.componentsIndex):
            for compName in self.componentsIndex:
                comp = self.components[compName]
                mgear.log(name + " : " + comp.fullName +
                          " (" + comp.type + ")")
                comp.stepMethods[i]()

            if self.options["step"] >= 1 and i >= self.options["step"] - 1:
                break

    def finalize(self):
        """Finalize the rig."""
        groupIdx = 0

        # Properties --------------------------------------
        mgear.log("Finalize")

        # clean jnt_org --------------------------------------
        if self.options["joint_rig"]:
            mgear.log("Cleaning jnt org")
            for jOrg in dag.findChildrenPartial(self.jnt_org, "org"):
                if not jOrg.listRelatives(c=True):
                    pm.delete(jOrg)

        # Groups ------------------------------------------
        mgear.log("Creating groups")
        # Retrieve group content from components
        for name in self.componentsIndex:
            component_ = self.components[name]
            for name, objects in component_.groups.items():
                self.addToGroup(objects, name)
            for name, objects in component_.subGroups.items():
                self.addToSubGroup(objects, name)

        # Create master set to group all the groups
        masterSet = pm.sets(n=self.model.name() + "_sets_grp", em=True)
        pm.connectAttr(masterSet.message, self.model.rigGroups[groupIdx])
        groupIdx += 1

        # Creating all groups
        pm.select(cl=True)
        for name, objects in self.groups.items():
            s = pm.sets(n=self.model.name() + "_" + name + "_grp")
            s.union(objects)
            pm.connectAttr(s.message, self.model.rigGroups[groupIdx])
            groupIdx += 1
            masterSet.add(s)
        for parentGroup, subgroups in self.subGroups.items():
            pg = pm.PyNode(self.model.name() + "_" + parentGroup + "_grp")
            for sg in subgroups:
                sub = pm.PyNode(self.model.name() + "_" + sg + "_grp")
                if sub in masterSet.members():
                    masterSet.remove(sub)
                pg.add(sub)

        # Bind pose ---------------------------------------
        # controls_grp = self.groups["controllers"]
        # pprint(controls_grp, stream=None, indent=1, width=100)
        ctl_master_grp = pm.PyNode(self.model.name() + "_controllers_grp")
        pm.select(ctl_master_grp, replace=True)
        node = pm.dagPose(save=True, selection=True)
        pm.connectAttr(node.message, self.model.rigPoses[0])
        print node

        # Bind skin re-apply
        if self.options["importSkin"]:
            try:
                pm.displayInfo("Importing Skin")
                skin.importSkin(self.options["skin"])

            except RuntimeError:
                pm.displayWarning(
                    "Skin doesn't exist or is not correct. " +
                    self.options["skin"] + " Skipped!")

    def addCtl(self, parent, name, m, color, iconShape, **kwargs):
        """Create the control and apply the shape, if this is alrealdy stored
        in the guide controllers grp.

        Args:
            parent (dagNode): The control parent
            name (str): The control name.
            m (matrix): The transfromation matrix for the control.
            color (int or list of float): The color for the control in index
                or RGB.
            iconShape (str): The controls default shape.
            kwargs (variant): Other arguments for the iconShape type variations

        Returns:
            dagNode: The Control.

        """
        if "degree" not in kwargs.keys():
            kwargs["degree"] = 1

        bufferName = name + "_controlBuffer"
        if bufferName in self.guide.controllers.keys():
            ctl_ref = self.guide.controllers[bufferName]
            ctl = primitive.addTransform(parent, name, m)
            for shape in ctl_ref.getShapes():
                ctl.addChild(shape, shape=True, add=True)
                pm.rename(shape, name + "Shape")
        else:
            ctl = icon.create(parent, name, m, color, iconShape, **kwargs)

        self.addToGroup(ctl, "controllers")

        # Set the control shapes isHistoricallyInteresting
        for oShape in ctl.getShapes():
            oShape.isHistoricallyInteresting.set(False)

        # set controller tag
        if versions.current() >= 201650:
            pm.controller(ctl)

        return ctl

    def addToGroup(self, objects, names=["hidden"]):
        """Add the object in a collection for later group creation.

        Args:
            objects (dagNode or list of dagNode): Object to put in the group.
            names (str or list of str): Names of the groups to create.

        """
        if not isinstance(names, list):
            names = [names]

        if not isinstance(objects, list):
            objects = [objects]

        for name in names:
            if name not in self.groups.keys():
                self.groups[name] = []

            self.groups[name].extend(objects)

    def addToSubGroup(self, subGroups, parentGroups=["hidden"]):
        """Add the object in a collection for later SubGroup creation.

        Args:
            subGroups (dagNode or list of dagNode): Groups (maya set) to add
                as a Subgroup.
            namparentGroupses (str or list of str): Names of the parent groups
                to create.

        """

        if not isinstance(parentGroups, list):
            parentGroups = [parentGroups]

        if not isinstance(subGroups, list):
            subGroups = [subGroups]

        for pg in parentGroups:
            if pg not in self.subGroups.keys():
                self.subGroups[pg] = []
            self.subGroups[pg].extend(subGroups)

    def getLocalName(self, guideName):
        """This function return the local name, cutting the Maya fullname
        and taking the latest part.

            ie. "parentA|parentB|arm_C0_root" will return "arm_C0_root"

        Args:
            guideName (str): The guide name.

        Returns:
            str: The local Name

        """
        if guideName is None:
            return None
        localName = guideName.split("|")[-1]
        return localName

    def getComponentName(self, guideName, local=True):
        """
        This function return the component name

            ie. "arm_C0_root" return "arm_C0"

        Args:
            guideName (str): The guide name.

        Returns:
            str: The compnent Name
        """

        if guideName is None:
            return None
        if local:
            guideName = self.getLocalName(guideName)

        comp_name = "_".join(guideName.split("_")[:2])
        return comp_name

    def getRelativeName(self, guideName):
        """This function return the name of the relative in the guide

            ie. "arm_C0_root" return "root"

        Args:
            guideName (str): The guide name.

        Returns:
            str: The relative Name

        """
        if guideName is None:
            return None

        localName = self.getLocalName(guideName)
        relative_name = "_".join(localName.split("_")[2:])
        return relative_name

    def findRelative(self, guideName):
        """Return the objects in the rig matching the guide object.

        Args:
            guideName (str): Name of the guide object.

        Returns:
           transform: The relative object

        """

        if guideName is None:
            return self.global_ctl

        # localName = self.getLocalName(guideName)
        comp_name = self.getComponentName(guideName)
        relative_name = self.getRelativeName(guideName)

        if comp_name not in self.components.keys():
            return self.global_ctl
        return self.components[comp_name].getRelation(relative_name)

    def findControlRelative(self, guideName):
        """Return the control objects in the rig matching the guide object.

        Args:
            guideName (str): Name of the guide object.

        Returns:
           transform: The relative control object

        """

        if guideName is None:
            return self.global_ctl

        # localName = self.getLocalName(guideName)
        comp_name = self.getComponentName(guideName)
        relative_name = self.getRelativeName(guideName)

        if comp_name not in self.components.keys():
            return self.global_ctl
        return self.components[comp_name].getControlRelation(relative_name)

    # TODO: update findComponent and other find methods with new funtions like
    # comp_name and others.  Better composability
    def findComponent(self, guideName):
        """Return the component from a guide Name.

        Args:
            guideName (str): Name of the guide object.

        Returns:
           transform: The component

        """
        if guideName is None:
            return None

        comp_name = self.getComponentName(guideName, False)
        # comp_name = "_".join(guideName.split("_")[:2])

        if comp_name not in self.components.keys():
            return None

        return self.components[comp_name]

    def findUIHost(self, guideName):
        """Return the UI host of the compoent

        Args:
            guideName (str): Name of the guide object.

        Returns:
           transform: The relative object

        """

        if guideName is None:
            return self.ui

        comp_name = self.getComponentName(guideName, False)
        # comp_name = "_".join(guideName.split("_")[:2])

        if comp_name not in self.components.keys():
            return self.ui

        if self.components[comp_name].ui is None:
            self.components[comp_name].ui = pm.UIHost(
                self.components[comp_name].root)

        return self.components[comp_name].ui
