import os
import tempfile
from typing import Optional

import FreeCADGui as Gui
import FreeCAD as App

class AHB_Render:
    def GetResources(self):
        return {"MenuText": "Render",
                "ToolTip": "Render to images",
                "Pixmap": ""
                }

    def IsActive(self):
        return True

    def set_render_lines(self, line_color = (0.0,0.0,0.0,0.0), background_color = (1.0,1.0,1.0,0.0), mask_stages_below: Optional[int] = None, mask_color=(1.0,1.0,1.0)):
        doc = App.activeDocument()
        for obj in doc.Objects:
            if obj.TypeId == 'Part::Feature':
                if 'AssemblyHandbook_Stage' in obj.PropertiesList:
                    masked = mask_stages_below is not None and obj.AssemblyHandbook_Stage < mask_stages_below
                    if 'AssemblyHandbook_RenderLines' in obj.PropertiesList and not obj.AssemblyHandbook_RenderLines:
                        obj.ViewObject.LineColor = background_color
                    else:
                        obj.ViewObject.LineColor = line_color
                    obj.ViewObject.DisplayMode = 'Flat Lines' if not masked else 'Shaded'
                    obj.ViewObject.ShapeMaterial.AmbientColor = (0.0, 0.0, 0.0, 0.0)
                    obj.ViewObject.ShapeMaterial.DiffuseColor = (0.0, 0.0, 0.0, 0.0)
                    obj.ViewObject.ShapeMaterial.SpecularColor = (0.0, 0.0, 0.0, 0.0)
                    obj.ViewObject.ShapeMaterial.EmissiveColor = background_color if not masked else mask_color

    def set_render_outlines(self, mask_stages_below: Optional[int] = None):
        doc = App.activeDocument()
        step = 8
        r = step
        g = step
        b = step
        for obj in doc.Objects:
            if obj.TypeId == 'Part::Feature':
                if 'AssemblyHandbook_Stage' in obj.PropertiesList:
                    masked = mask_stages_below is not None and obj.AssemblyHandbook_Stage < mask_stages_below
                    obj.ViewObject.DisplayMode = 'Shaded'
                    obj.ViewObject.ShapeMaterial.AmbientColor = (0.0, 0.0, 0.0, 0.0)
                    obj.ViewObject.ShapeMaterial.DiffuseColor = (0.0, 0.0, 0.0, 0.0)
                    obj.ViewObject.ShapeMaterial.SpecularColor = (0.0, 0.0, 0.0, 0.0)
                    #obj.ViewObject.ShapeMaterial.EmissiveColor = (r/255.0, g/255.0, b/255.0, 0.0) if not masked else (1.0,1.0,1.0,0.0)
                    obj.ViewObject.ShapeMaterial.EmissiveColor = (r/255.0, g/255.0, b/255.0, 0.0)
                    if masked:
                        obj.ViewObject.Visibility = False

                    r = r + step
                    if r >= 256 - step:
                        r = step
                        g = g + step
                        if g >= 256 - step:
                            g = step
                            b = b + step
                            if b >= 256 - step:
                                b = step

    def reset_display(self):
        doc = App.activeDocument()
        for obj in doc.Objects:
            if obj.TypeId == 'Part::Feature':
                if 'AssemblyHandbook_Stage' in obj.PropertiesList:
                    obj.ViewObject.LineColor = (0.0, 0.0, 0.0, 0.0)
                    obj.ViewObject.DisplayMode = 'Flat Lines'
                    obj.ViewObject.ShapeMaterial.AmbientColor = (0.3, 0.3, 0.3, 0.0)
                    obj.ViewObject.ShapeMaterial.DiffuseColor = (1.0, 1.0, 1.0, 0.0)
                    obj.ViewObject.ShapeMaterial.SpecularColor = (0.5, 0.5, 0.5, 0.0)
                    obj.ViewObject.ShapeMaterial.EmissiveColor = (0.0, 0.0, 0.0, 0.0)

        workbench = Gui.getWorkbench("AssemblyHandbookWorkbench")
        workbench.context.onPartStageChanged(None)

    def render(self, resolution, filename: str, line_color=(0.0, 0.0, 0.0, 0.0), fill_color=(1.0, 1.0, 1.0, 0.0), mask_stages_below: Optional[int] = None):
        import time
        from PIL import Image, ImageFilter

        render_start_time = time.perf_counter()

        temp_lines_file_name = tempfile.gettempdir() + "/ahb_temp_lines.png"
        temp_shapes_file_name = tempfile.gettempdir() + "/ahb_temp_shapes.png"

        #temp_lines_file_name = filename + "-lines.png"
        #temp_shapes_file_name = filename + "-shapes.png"

        # render lines in black, background in red, fill shapes in green
        # the green band contains the lines images, the red band contains the inverted alpha layer
        self.set_render_lines((0.0, 0.0, 0.0), (0.0, 1.0, 0.0), mask_stages_below=mask_stages_below, mask_color=(1.0, 0.0, 1.0))
        Gui.ActiveDocument.ActiveView.saveImage(temp_lines_file_name, resolution[0], resolution[1], "#ff0000")

        self.set_render_outlines(mask_stages_below=mask_stages_below)
        Gui.ActiveDocument.ActiveView.saveImage(temp_shapes_file_name, resolution[0] * 2, resolution[1] * 2, "#ffffff")

        self.reset_display()

        lines_bands = Image.open(temp_lines_file_name).split()
        lines = lines_bands[1]
        alpha_band = lines_bands[0].point(lambda p: 255 - p)
        shapes = Image.open(temp_shapes_file_name)

        outlines_start_time = time.perf_counter()

        outlines = None
        for x in range(0, 3):
            for y in range(0, 3):
                if x == 1 and y == 1: continue
                kernel = [0, 0, 0, 0, 1, 0, 0, 0, 0]
                kernel[y * 3 + x] = -1
                partial_outlines = shapes.filter(ImageFilter.Kernel((3, 3), kernel, 1, 127))
                partial_outlines = partial_outlines.point(lambda p: 255 if p == 127 else 0)
                partial_outlines = partial_outlines.convert("L")
                partial_outlines = partial_outlines.point(lambda p: 255 if p == 255 else 0)
                if outlines is None:
                    outlines = partial_outlines
                else:
                    outlines.paste(partial_outlines, None, partial_outlines.point(lambda p: 0 if p == 255 else 255))

        # erase masked outlines
        outlines.paste(outlines.point(lambda p: 255), None, lines_bands[2].resize(outlines.size).point(lambda p: 255 if p == 255 else 0))

        # outlines.save("/home/youen/dev_linux/vhelio-render/vhelio-outlines.png")

        # outlines = outlines.resize(lines.size, Image.BILINEAR)
        # lines.paste(outlines, None, outlines.point(lambda p: 255 - p))

        lines_fullres = lines.resize(outlines.size, Image.NEAREST)
        lines_fullres.paste(outlines, None, outlines.point(lambda p: 255 if p == 0 else 0))
        lines = lines_fullres.resize(lines.size, Image.BILINEAR)

        alpha_band_fullres = alpha_band.resize(outlines.size, Image.NEAREST)
        alpha_band_fullres.paste(outlines.point(lambda p: 255), None, outlines.point(lambda p: 255 if p == 0 else 0))
        alpha_band = alpha_band_fullres.resize(lines.size, Image.BILINEAR)

        outlines_end_time = time.perf_counter()

        # colorize

        result = Image.merge("RGBA", [
            lines.point(lambda p: int(fill_color[0] * p + line_color[0]*(255.0-p))),
            lines.point(lambda p: int(fill_color[1] * p + line_color[1] * (255.0 - p))),
            lines.point(lambda p: int(fill_color[2] * p + line_color[2] * (255.0 - p))),
            alpha_band
        ])

        result.save(filename)

        print("Rendered " + filename + " in " + str(
            round((outlines_end_time - render_start_time) * 1000) / 1000) + "s (outlines detection in " + str(
            round((outlines_end_time - outlines_start_time) * 1000) / 1000) + "s)")

    def Activated(self):
        import shutil
        from PIL import Image, ImageFilter

        render_main = False
        render_stages = False
        render_parts = True

        workbench = Gui.getWorkbench("AssemblyHandbookWorkbench")
        doc = App.activeDocument()
        doc_file_name: str = doc.FileName
        if doc_file_name is None:
            raise BaseException("You must save your FreeCAD document before rendering images")

        filename = os.path.splitext(doc_file_name)[0] + ".png"
        dir = os.path.dirname(filename)

        if render_main:
            resolution = (2000,2000)
            workbench.context.setAllStagesVisible(True)
            self.render(resolution, filename)
            img_full = Image.new('RGB', resolution, (255, 255, 255))
            img = Image.open(filename)
            img_full.paste(img, None, img.getchannel('A'))
            img_full.save(filename)

        if render_stages:
            shutil.rmtree(dir + "/stages", ignore_errors=True)
            os.makedirs(dir + "/stages", exist_ok=True)
            all_stages = workbench.context.getAllStages()
            prev_stage_id: Optional[int] = None
            for stage_id in all_stages:
                stage_name = str(stage_id)
                while len(stage_name) < 6:
                    stage_name = "0" + stage_name
                resolution = (2000, 2000)
                if prev_stage_id is not None:
                    workbench.context.setActiveStage(prev_stage_id)
                    workbench.context.setAllStagesVisible(False)
                    self.render(resolution, dir + "/stages/" + stage_name + "-bg.png", (0.7, 0.7, 0.7), (1,1,1))
                workbench.context.setActiveStage(stage_id)
                workbench.context.setAllStagesVisible(False)
                self.render(resolution, dir + "/stages/" + stage_name + ".png", mask_stages_below=stage_id)

                # merge images
                bg = Image.new('RGB', resolution, (255,255,255))
                if prev_stage_id is not None:
                    prev_stage = Image.open(dir + "/stages/" + stage_name + "-bg.png")
                    prev_stage.load()
                    os.remove(dir + "/stages/" + stage_name + "-bg.png")
                    bg.paste(prev_stage, None, prev_stage.getchannel('A'))
                fg = Image.open(dir + "/stages/" + stage_name + ".png")
                bg.paste(fg, None, fg.getchannel('A'))
                bg.save(dir + "/stages/" + stage_name + ".png")

                if prev_stage_id is not None:
                    pass
                prev_stage_id = stage_id

        if render_parts:
            rendered_references = []
            resolution = (500, 500)
            shutil.rmtree(dir + "/parts", ignore_errors=True)
            os.makedirs(dir + "/parts", exist_ok=True)
            count = 0
            for part_to_render in doc.Objects:
                if part_to_render.TypeId == 'Part::Feature' and 'AssemblyHandbook_Stage' in part_to_render.PropertiesList:
                    if 'Base_Reference' not in part_to_render.PropertiesList or part_to_render.Base_Reference == "" or part_to_render.Base_Reference in rendered_references:
                        continue

                    rendered_references.append(part_to_render.Base_Reference)

                    for part_to_hide in doc.Objects:
                        if part_to_hide.TypeId == 'Part::Feature':
                            part_to_hide.ViewObject.Visibility = False
                    part_to_render.ViewObject.Visibility = True

                    Gui.ActiveDocument.ActiveView.fitAll()

                    part_filename = dir + "/parts/" + part_to_render.Base_Reference.upper().replace(' ', '-') + ".png"
                    self.render(resolution, part_filename)

                    bg = Image.new('RGB', resolution, (255, 255, 255))
                    fg = Image.open(part_filename)
                    bg.paste(fg, None, fg.getchannel('A'))
                    bg.save(part_filename)

                    count = count + 1
                    if count == 10:
                        pass

        workbench.context.setAllStagesVisible(True)

from ahb_command import AHB_CommandWrapper
AHB_CommandWrapper.addGuiCommand('AHB_render', AHB_Render())