import FreeCADGui as Gui
import FreeCAD as App
import random

def _rgb(r, g, b):
    return r/255.0, g/255.0, b/255.0, 0.0

_colors = [
    _rgb(94, 224, 174),
    _rgb(223, 234, 70),
    _rgb(35, 120, 200),
    _rgb(188, 80, 33),
    _rgb(26, 242, 69),
    _rgb(13, 41, 107),
    _rgb(153, 9, 52),
    _rgb(211, 119, 0),
    _rgb(34, 130, 7),
    _rgb(221, 79, 249)
]

class AHB_Context:
    _activeStageId: int = None

    def setActiveStage(self, stageId):
        self._activeStageId = stageId

    def getActiveStage(self):
        return self._activeStageId

    def onPartStageChanged(self, part: App.DocumentObject):
        allStages = set()

        doc = App.ActiveDocument
        for obj in doc.Objects:
            if obj.TypeId == 'Part::Feature':
                if 'AssemblyHandbook_Stage' in obj.PropertiesList:
                    allStages.add(obj.AssemblyHandbook_Stage)

        allStages = list(allStages)
        allStages.sort()

        if not 'AssemblyHandbook_Stage' in part.PropertiesList:
            part.ViewObject.ShapeColor = (0.8,0.8,0.8,0.0)

        for obj in doc.Objects:
            if obj.TypeId == 'Part::Feature':
                if 'AssemblyHandbook_Stage' in obj.PropertiesList:
                    stageIndex = allStages.index(obj.AssemblyHandbook_Stage)
                    obj.ViewObject.ShapeColor = _colors[stageIndex % len(_colors)]