from OCC.Core.Quantity import Quantity_Color, Quantity_NOC_PERU, Quantity_ColorRGBA
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform, BRepBuilderAPI_GTransform
from OCC.Core.gp import gp_Trsf
from OCC.Display.OCCViewer import rgb_color
from OCC.Core.AIS import AIS_Shape
from OCC.Core.TopExp import TopExp_Explorer
from OCC.Core.TopAbs import TopAbs_FACE
from OCC.Core.BRep import BRep_Builder
from OCC.Core.TopoDS import TopoDS_Compound

from bcad.binterpreter.singleton import Singleton
from logging import debug

DISP_MODE_NONE = 0
DISP_MODE_SHADED = 1
DISP_MODE_WIREFRAME = 2


class SCLShape(object):
    def __init__(self, shape):
        self.trsf = gp_Trsf()
        self.shape = shape
        self.shape_rgba: Quantity_ColorRGBA | None = None
        self._face_colors: dict[int, Quantity_ColorRGBA] | None = None
        self.display_mode = DISP_MODE_NONE
        self.style = "main"
        self.hidden = False
        self.context_role = ""

    def set_hidden(self, hidden):
        self.hidden = hidden

    def is_hidden(self):
        return self.hidden

    def set_linestyle(self, name):
        if (name == "hidden"):
            debug("Set hidden line style")
            self.style = "hidden"
        if (name == "main_projection"):
            debug("Set main_projection line style")
            self.style = "main_projection"

    def set_shape_color(self, shape_rgba):
        self.shape_rgba = shape_rgba

    def get_shape_color(self):
        return self.shape_rgba

    def get_default_color(self) -> Quantity_ColorRGBA:
        if self.shape_rgba is not None:
            return self.shape_rgba
        if self.context_role == "cutter":
            return Quantity_ColorRGBA(rgb_color(0.8, 0.2, 0.2), 1.0)
        return Quantity_ColorRGBA(Quantity_Color(Quantity_NOC_PERU), 1.0)

    def set_face_colors(self, face_colors: dict[int, Quantity_ColorRGBA]):
        self._face_colors = face_colors

    def get_face_colors(self):
        return self._face_colors

    def has_face_colors(self):
        return self._face_colors is not None and len(self._face_colors) > 0

    def color(self, color):
        debug("Trying to set color: %s"%(color,))
        if color is None:
            self.shape_rgba = None
            self._face_colors = None
            return

        if isinstance(color, list) and len(color) >= 4:
            self.shape_rgba = Quantity_ColorRGBA(
                float(color[0]), float(color[1]), float(color[2]),
                float(color[3]))
            self._face_colors = None

    def get_shape(self):
        return self.shape

    def transform(self, trsf):
        self.trsf = trsf
        self.transform_shape()

    def transform_shape(self):
        if self.shape is not None and self._face_colors:
            old_hashes = []
            exp = TopExp_Explorer(self.shape, TopAbs_FACE)
            while exp.More():
                old_hashes.append(hash(exp.Current()))
                exp.Next()

        self.shape = BRepBuilderAPI_Transform(self.shape, self.trsf, True).Shape()

        if self._face_colors and old_hashes:
            new_colors = {}
            exp = TopExp_Explorer(self.shape, TopAbs_FACE)
            for old_hash in old_hashes:
                if exp.More():
                    color = self._face_colors.get(old_hash)
                    if color is not None:
                        new_colors[hash(exp.Current())] = color
                    exp.Next()
            self._face_colors = new_colors

    def transform_gtrsf(self, gtrsf):
        if self.shape is not None and self._face_colors:
            old_hashes = []
            exp = TopExp_Explorer(self.shape, TopAbs_FACE)
            while exp.More():
                old_hashes.append(hash(exp.Current()))
                exp.Next()

        self.shape = BRepBuilderAPI_GTransform(self.shape, gtrsf, True).Shape()

        if self._face_colors and old_hashes:
            new_colors = {}
            exp = TopExp_Explorer(self.shape, TopAbs_FACE)
            for old_hash in old_hashes:
                if exp.More():
                    color = self._face_colors.get(old_hash)
                    if color is not None:
                        new_colors[hash(exp.Current())] = color
                    exp.Next()
            self._face_colors = new_colors

    def display_hidden(self, ais_context):
        ais_shp = AIS_Shape(self.shape)
        ais_shp.SetWidth(0.1)
        ais_shp.SetTransparency(0.10)
        ais_shp.SetColor(rgb_color(0,0,0))
        aspect = ais_shp.Attributes().WireAspect()
        aspect.SetColor(rgb_color(0,0,0))
        aspect.SetTypeOfLine(1)
        ais_context.Display(ais_shp, True)

    def display_main_projection(self, ais_context):
        ais_shp = AIS_Shape(self.shape)
        ais_shp.SetWidth(5)
        aspect = ais_shp.Attributes().WireAspect()
        aspect.SetColor(rgb_color(0,0,0))
        aspect.SetTypeOfLine(0)
        ais_context.Display(ais_shp, True)

    def display_main(self, ais_context):
        if self.has_face_colors():
            self._display_main_face_colors(ais_context)
            return

        ais_shp = AIS_Shape(self.shape)
        ais_shp.SetWidth(2.0)
        rgba = self.get_default_color()
        ais_shp.SetColor(rgb_color(rgba.GetRGB().Red(), rgba.GetRGB().Green(), rgba.GetRGB().Blue()))
        if rgba.Alpha() < 1.0:
            ais_shp.SetTransparency(1.0 - rgba.Alpha())
        ais_context.Display(ais_shp, False)

    def _display_main_face_colors(self, ais_context):
        color_to_faces = {}
        exp = TopExp_Explorer(self.shape, TopAbs_FACE)
        while exp.More():
            face = exp.Current()
            color = self._face_colors.get(hash(face), self.get_default_color())
            key = (color.GetRGB().Red(), color.GetRGB().Green(), color.GetRGB().Blue(), color.Alpha())
            if key not in color_to_faces:
                color_to_faces[key] = (color, [])
            color_to_faces[key][1].append(face)
            exp.Next()

        for color, faces in color_to_faces.values():
            compound = TopoDS_Compound()
            BRep_Builder().MakeCompound(compound)
            for face in faces:
                BRep_Builder().Add(compound, face)

            ais_shp = AIS_Shape(compound)
            ais_shp.SetWidth(2.0)
            ais_shp.SetColor(rgb_color(color.GetRGB().Red(), color.GetRGB().Green(), color.GetRGB().Blue()))
            if color.Alpha() < 1.0:
                ais_shp.SetTransparency(1.0 - color.Alpha())
            ais_context.Display(ais_shp, False)


    def display(self, writer=None):
        #debug("shape Display")
        if self.shape is not None and not self.shape.IsNull():
            if (writer is not None):
                if hasattr(writer, 'TransferColor'):
                    if self.has_face_colors():
                        self._write_face_colors(writer)
                    else:
                        writer.TransferColor(self.shape, self.shape_rgba)
                else:
                    writer.Transfer(self.shape)
            else:
                ais_context = Singleton.sd.display.GetContext()
                if (self.style == 'hidden'):
                    self.display_hidden(ais_context)
                elif (self.style == 'main_projection'):
                    self.display_main_projection(ais_context)
                else:
                    self.display_main(ais_context)
        else:
            #warning("Empty shape")
            pass

    def _write_face_colors(self, writer):
        color_to_faces = {}
        exp = TopExp_Explorer(self.shape, TopAbs_FACE)
        while exp.More():
            face = exp.Current()
            color = self._face_colors.get(hash(face), self.get_default_color())
            key = (color.GetRGB().Red(), color.GetRGB().Green(), color.GetRGB().Blue(), color.Alpha())
            if key not in color_to_faces:
                color_to_faces[key] = (color, [])
            color_to_faces[key][1].append(face)
            exp.Next()

        for color, faces in color_to_faces.values():
            compound = TopoDS_Compound()
            BRep_Builder().MakeCompound(compound)
            for face in faces:
                BRep_Builder().Add(compound, face)
            writer.TransferColor(compound, color)
