from __future__ import annotations

import colorsys
import json
from pathlib import Path
from typing import Any

import bpy
from bpy.props import (
    BoolProperty,
    CollectionProperty,
    EnumProperty,
    FloatVectorProperty,
    IntProperty,
    PointerProperty,
    StringProperty,
)
from bpy.types import Operator, Panel, PropertyGroup


_TEMP_PREFIX = "__MID_BAKER_"


def _distinct_color(index: int) -> tuple[float, float, float]:
    """Return a deterministic, vivid linear RGB color for a zero-based slot."""
    hue = (0.03 + index * 0.618033988749895) % 1.0
    saturation = 0.72 if index % 2 == 0 else 0.9
    value = 1.0 if index % 3 else 0.82
    return colorsys.hsv_to_rgb(hue, saturation, value)


def _index_color(index: int) -> tuple[float, float, float]:
    """Encode a one-based material ID into 24 bits, reserving black for empty pixels."""
    value = index + 1
    return (
        (value & 0xFF) / 255.0,
        ((value >> 8) & 0xFF) / 255.0,
        ((value >> 16) & 0xFF) / 255.0,
    )


def _slot_color(
    index: int,
    material: bpy.types.Material | None,
    palette_mode: str,
) -> tuple[float, float, float]:
    if palette_mode == "MATERIAL" and material is not None:
        return tuple(material.diffuse_color[:3])
    if palette_mode == "INDEX":
        return _index_color(index)
    return _distinct_color(index)


def _create_id_material(
    slot_index: int,
    color: tuple[float, float, float],
    image: bpy.types.Image,
) -> bpy.types.Material:
    material = bpy.data.materials.new(f"{_TEMP_PREFIX}Material_{slot_index + 1}")
    material.use_nodes = True
    nodes = material.node_tree.nodes
    nodes.clear()

    output = nodes.new("ShaderNodeOutputMaterial")
    output.location = (360.0, 0.0)
    emission = nodes.new("ShaderNodeEmission")
    emission.location = (80.0, 0.0)
    emission.inputs["Color"].default_value = (*color, 1.0)
    emission.inputs["Strength"].default_value = 1.0

    image_node = nodes.new("ShaderNodeTexImage")
    image_node.location = (-260.0, -180.0)
    image_node.image = image
    image_node.select = True
    nodes.active = image_node

    material.node_tree.links.new(emission.outputs["Emission"], output.inputs["Surface"])
    return material


def _create_baked_ids_material(image: bpy.types.Image) -> bpy.types.Material:
    """Create an export-friendly, single material displaying the baked ID image."""
    material = bpy.data.materials.new("baked ids")
    material.use_nodes = True
    nodes = material.node_tree.nodes
    nodes.clear()

    output = nodes.new("ShaderNodeOutputMaterial")
    output.location = (560.0, 0.0)
    principled = nodes.new("ShaderNodeBsdfPrincipled")
    principled.location = (260.0, 0.0)
    principled.inputs["Roughness"].default_value = 1.0

    image_node = nodes.new("ShaderNodeTexImage")
    image_node.label = "Baked Material IDs"
    image_node.location = (-120.0, 0.0)
    image_node.image = image
    image_node.interpolation = "Closest"
    nodes.active = image_node

    material.node_tree.links.new(image_node.outputs["Color"], principled.inputs["Base Color"])
    material.node_tree.links.new(principled.outputs["BSDF"], output.inputs["Surface"])
    return material


def _create_export_copy(
    source: bpy.types.Object,
    baked_mesh: bpy.types.Mesh,
    image: bpy.types.Image,
    apply_modifiers: bool,
    target_collection: bpy.types.Collection,
) -> tuple[bpy.types.Object, bpy.types.Mesh, bpy.types.Material]:
    """Create a persistent one-material copy suitable for later export."""
    if apply_modifiers:
        export_mesh = baked_mesh.copy()
        export_object = bpy.data.objects.new(
            f"{source.name}_Baked_IDs",
            export_mesh,
        )
        export_object.matrix_world = source.matrix_world.copy()
    else:
        export_object = source.copy()
        export_mesh = source.data.copy()
        export_object.data = export_mesh
        export_object.name = f"{source.name}_Baked_IDs"

    target_collection.objects.link(export_object)
    export_object.hide_set(False)
    export_object.hide_viewport = False
    export_object.hide_render = False

    material = _create_baked_ids_material(image)
    export_mesh.materials.clear()
    export_mesh.materials.append(material)
    for polygon in export_mesh.polygons:
        polygon.material_index = 0

    return export_object, export_mesh, material


def _normalise_png_path(filepath: str) -> Path:
    path = Path(bpy.path.abspath(filepath)).expanduser()
    if path.suffix.lower() != ".png":
        path = path.with_suffix(".png")
    return path


def _normalise_glb_path(filepath: str) -> Path:
    path_text = filepath.strip()
    if not path_text:
        raise RuntimeError("Choose a GLB export path")
    path = Path(bpy.path.abspath(path_text)).expanduser()
    if path.suffix.lower() != ".glb":
        path = path.with_suffix(".glb")
    return path


def _quick_export_glb(
    context: bpy.types.Context,
    export_object: bpy.types.Object,
    filepath: str,
) -> Path:
    output_path = _normalise_glb_path(filepath)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    for obj in list(context.selected_objects):
        obj.select_set(False)
    export_object.select_set(True)
    context.view_layer.objects.active = export_object

    result = bpy.ops.export_scene.gltf(
        filepath=str(output_path),
        export_format="GLB",
        use_selection=True,
        export_materials="EXPORT",
        export_animations=False,
    )
    if "FINISHED" not in result:
        raise RuntimeError("Blender did not finish the GLB export")
    return output_path


def _snapshot_bake_settings(scene: bpy.types.Scene) -> dict[str, Any]:
    bake = scene.render.bake
    return {
        "engine": scene.render.engine,
        "margin": bake.margin,
        "target": bake.target,
        "use_clear": bake.use_clear,
        "use_selected_to_active": bake.use_selected_to_active,
        "cycles_samples": scene.cycles.samples,
    }


def _restore_bake_settings(scene: bpy.types.Scene, state: dict[str, Any]) -> None:
    scene.render.engine = state["engine"]
    bake = scene.render.bake
    bake.margin = state["margin"]
    bake.target = state["target"]
    bake.use_clear = state["use_clear"]
    bake.use_selected_to_active = state["use_selected_to_active"]
    scene.cycles.samples = state["cycles_samples"]


class MIDB_PaletteEntry(PropertyGroup):
    slot: IntProperty(name="Slot", default=1, min=1)
    material_name: StringProperty(name="Material")
    color: FloatVectorProperty(
        name="ID Color",
        subtype="COLOR",
        size=3,
        min=0.0,
        max=1.0,
    )


class MIDB_Settings(PropertyGroup):
    resolution: EnumProperty(
        name="Resolution",
        items=(
            ("256", "256", "256 x 256"),
            ("512", "512", "512 x 512"),
            ("1024", "1K", "1024 x 1024"),
            ("2048", "2K", "2048 x 2048"),
            ("4096", "4K", "4096 x 4096"),
            ("8192", "8K", "8192 x 8192"),
        ),
        default="2048",
    )
    margin: IntProperty(
        name="Margin",
        description="Extend IDs beyond UV island edges by this many pixels",
        default=1,
        min=0,
        max=256,
        subtype="PIXEL",
    )
    palette_mode: EnumProperty(
        name="Colors",
        items=(
            (
                "DISTINCT",
                "Distinct",
                "Generate deterministic, easy-to-see colors for material slots",
            ),
            (
                "MATERIAL",
                "Material Viewport Colors",
                "Use each material's viewport display color",
            ),
            (
                "INDEX",
                "Exact Slot Index",
                "Encode the one-based slot number as a 24-bit RGB integer",
            ),
        ),
        default="DISTINCT",
    )
    apply_modifiers: BoolProperty(
        name="Apply Modifiers",
        description="Bake evaluated modifier geometry instead of the base mesh",
        default=False,
    )
    create_export_copy: BoolProperty(
        name="Create Export Copy",
        description=(
            "Create and select a duplicate with one 'baked ids' material using the baked image"
        ),
        default=False,
    )
    quick_export_glb: BoolProperty(
        name="Quick Export GLB",
        description="Immediately export the generated one-material copy as a GLB",
        default=False,
    )
    glb_filepath: StringProperty(
        name="GLB Path",
        description="Path for the Substance Painter-ready GLB export",
        subtype="FILE_PATH",
        default="//baked_ids.glb",
    )
    image_name: StringProperty(
        name="Image Name",
        default="Material_ID",
    )
    save_to_disk: BoolProperty(
        name="Save PNG",
        default=True,
    )
    filepath: StringProperty(
        name="File Path",
        subtype="FILE_PATH",
        default="//material_id.png",
    )
    write_legend: BoolProperty(
        name="Save JSON Legend",
        description="Write slot names and exact colors beside the PNG",
        default=True,
    )
    last_image_name: StringProperty(name="Last Image")
    last_palette: CollectionProperty(type=MIDB_PaletteEntry)


class MIDB_OT_Bake(Operator):
    bl_idname = "object.bake_material_id"
    bl_label = "Bake Material ID"
    bl_description = "Bake the active mesh's material assignments into an ID image"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context: bpy.types.Context) -> bool:
        obj = context.active_object
        return bool(
            obj
            and obj.type == "MESH"
            and context.mode == "OBJECT"
            and obj.data.uv_layers
        )

    def execute(self, context: bpy.types.Context) -> set[str]:
        settings = context.scene.material_id_baker
        source = context.active_object
        scene = context.scene
        view_layer = context.view_layer

        if source is None or source.type != "MESH":
            self.report({"ERROR"}, "Select an active mesh object")
            return {"CANCELLED"}
        if context.mode != "OBJECT":
            self.report({"ERROR"}, "Switch to Object Mode before baking")
            return {"CANCELLED"}
        if not source.data.uv_layers:
            self.report({"ERROR"}, "The active mesh needs a UV map")
            return {"CANCELLED"}

        selected_before = list(context.selected_objects)
        active_before = view_layer.objects.active
        bake_state = _snapshot_bake_settings(scene)

        temp_object: bpy.types.Object | None = None
        temp_mesh: bpy.types.Mesh | None = None
        temp_materials: list[bpy.types.Material] = []
        image: bpy.types.Image | None = None
        export_object: bpy.types.Object | None = None
        export_mesh: bpy.types.Mesh | None = None
        export_material: bpy.types.Material | None = None
        succeeded = False

        try:
            if settings.apply_modifiers:
                depsgraph = context.evaluated_depsgraph_get()
                evaluated = source.evaluated_get(depsgraph)
                temp_mesh = bpy.data.meshes.new_from_object(
                    evaluated,
                    preserve_all_data_layers=True,
                    depsgraph=depsgraph,
                )
            else:
                temp_mesh = source.data.copy()

            if not temp_mesh.uv_layers:
                raise RuntimeError("The evaluated mesh has no UV map")

            source_uv = source.data.uv_layers.active
            baked_uv = temp_mesh.uv_layers.get(source_uv.name) if source_uv else None
            if baked_uv is None:
                baked_uv = temp_mesh.uv_layers.active
            temp_mesh.uv_layers.active = baked_uv
            for uv_layer in temp_mesh.uv_layers:
                uv_layer.active_render = uv_layer == baked_uv

            temp_object = bpy.data.objects.new(
                f"{_TEMP_PREFIX}{source.name}",
                temp_mesh,
            )
            temp_object.matrix_world = source.matrix_world.copy()
            scene.collection.objects.link(temp_object)

            highest_slot = max(
                (polygon.material_index for polygon in temp_mesh.polygons),
                default=0,
            )
            slot_count = max(1, len(temp_mesh.materials), highest_slot + 1)
            source_materials = []
            for index in range(slot_count):
                material = None
                if index < len(source.material_slots):
                    material = source.material_slots[index].material
                if material is None and index < len(temp_mesh.materials):
                    material = temp_mesh.materials[index]
                source_materials.append(material)
            colors = [
                _slot_color(index, material, settings.palette_mode)
                for index, material in enumerate(source_materials)
            ]

            resolution = int(settings.resolution)
            image = bpy.data.images.new(
                name=settings.image_name.strip() or "Material_ID",
                width=resolution,
                height=resolution,
                alpha=True,
                float_buffer=False,
            )
            image.generated_color = (0.0, 0.0, 0.0, 0.0)
            image.alpha_mode = "STRAIGHT"
            try:
                image.colorspace_settings.name = "Non-Color"
            except TypeError:
                # Some color configurations use a differently named raw space.
                image.colorspace_settings.is_data = True

            for index, color in enumerate(colors):
                temp_materials.append(_create_id_material(index, color, image))

            polygon_material_indices = [
                polygon.material_index for polygon in temp_mesh.polygons
            ]
            temp_mesh.materials.clear()
            for material in temp_materials:
                temp_mesh.materials.append(material)
            for polygon, material_index in zip(
                temp_mesh.polygons,
                polygon_material_indices,
                strict=True,
            ):
                polygon.material_index = min(material_index, slot_count - 1)

            for obj in list(context.selected_objects):
                obj.select_set(False)
            temp_object.hide_set(False)
            temp_object.hide_render = False
            temp_object.select_set(True)
            view_layer.objects.active = temp_object

            # Cycles owns Blender's image-bake pipeline even though this flat
            # emission bake needs only one deterministic sample.
            scene.render.engine = "CYCLES"
            scene.cycles.samples = 1
            scene.render.bake.margin = settings.margin
            scene.render.bake.target = "IMAGE_TEXTURES"
            scene.render.bake.use_clear = True
            scene.render.bake.use_selected_to_active = False

            context.window_manager.progress_begin(0, 1)
            try:
                result = bpy.ops.object.bake(type="EMIT")
            finally:
                context.window_manager.progress_end()
            if "FINISHED" not in result:
                raise RuntimeError("Blender did not finish the bake")

            palette = [
                {
                    "slot": index + 1,
                    "material": material.name if material else f"Unassigned Slot {index + 1}",
                    "rgb": [round(component, 8) for component in colors[index]],
                }
                for index, material in enumerate(source_materials)
            ]
            image["material_id_palette"] = json.dumps(palette, separators=(",", ":"))
            image["source_object"] = source.name
            image["uv_map"] = temp_mesh.uv_layers.active.name

            saved_path: Path | None = None
            if settings.save_to_disk:
                saved_path = _normalise_png_path(settings.filepath)
                saved_path.parent.mkdir(parents=True, exist_ok=True)
                image.filepath_raw = str(saved_path)
                image.file_format = "PNG"
                image.save()

                if settings.write_legend:
                    legend_path = saved_path.with_suffix(".json")
                    legend = {
                        "image": saved_path.name,
                        "source_object": source.name,
                        "uv_map": temp_mesh.uv_layers.active.name,
                        "palette_mode": settings.palette_mode,
                        "materials": palette,
                    }
                    with legend_path.open("w", encoding="utf-8") as handle:
                        json.dump(legend, handle, indent=2)
                        handle.write("\n")

            settings.last_image_name = image.name
            settings.last_palette.clear()
            for item in palette:
                entry = settings.last_palette.add()
                entry.slot = item["slot"]
                entry.material_name = item["material"]
                entry.color = item["rgb"]

            if settings.create_export_copy:
                export_object, export_mesh, export_material = _create_export_copy(
                    source,
                    temp_mesh,
                    image,
                    settings.apply_modifiers,
                    context.collection,
                )

            exported_glb_path: Path | None = None
            if export_object is not None and settings.quick_export_glb:
                exported_glb_path = _quick_export_glb(
                    context,
                    export_object,
                    settings.glb_filepath,
                )

            succeeded = True
            if saved_path:
                message = f"Baked material IDs to {saved_path}"
            else:
                message = f"Baked material IDs to image '{image.name}'"
            if export_object is not None:
                message += f" and created '{export_object.name}'"
            if exported_glb_path is not None:
                message += f"; exported GLB to {exported_glb_path}"
            self.report({"INFO"}, message)
            return {"FINISHED"}

        except Exception as exc:
            self.report({"ERROR"}, f"Material ID bake failed: {exc}")
            return {"CANCELLED"}

        finally:
            if temp_object is not None:
                bpy.data.objects.remove(temp_object, do_unlink=True)
            if temp_mesh is not None and temp_mesh.users == 0:
                bpy.data.meshes.remove(temp_mesh)
            for material in temp_materials:
                if material.users == 0:
                    bpy.data.materials.remove(material)

            if not succeeded and export_object is not None:
                bpy.data.objects.remove(export_object, do_unlink=True)
            if not succeeded and export_mesh is not None and export_mesh.users == 0:
                bpy.data.meshes.remove(export_mesh)
            if not succeeded and export_material is not None and export_material.users == 0:
                bpy.data.materials.remove(export_material)

            _restore_bake_settings(scene, bake_state)

            for obj in list(context.selected_objects):
                obj.select_set(False)
            for obj in selected_before:
                if obj.name in view_layer.objects:
                    obj.select_set(True)
            if active_before and active_before.name in view_layer.objects:
                view_layer.objects.active = active_before

            if succeeded and export_object is not None:
                for obj in list(context.selected_objects):
                    obj.select_set(False)
                export_object.select_set(True)
                view_layer.objects.active = export_object

            if not succeeded and image is not None and image.users == 0:
                bpy.data.images.remove(image)


class MIDB_PT_Panel(Panel):
    bl_label = "Material ID Baker"
    bl_idname = "MIDB_PT_material_id_baker"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_category = "Material ID"

    def draw(self, context: bpy.types.Context) -> None:
        layout = self.layout
        settings = context.scene.material_id_baker
        obj = context.active_object

        if obj is None or obj.type != "MESH":
            layout.label(text="Select a mesh object", icon="INFO")
            return

        mesh = obj.data
        if mesh.uv_layers:
            uv_layer = mesh.uv_layers.active
            layout.label(text=f"UV Map: {uv_layer.name}", icon="GROUP_UVS")
        else:
            layout.label(text="A UV map is required", icon="ERROR")

        layout.label(text=f"Material Slots: {max(1, len(obj.material_slots))}")

        column = layout.column(align=True)
        column.prop(settings, "resolution")
        column.prop(settings, "margin")
        column.prop(settings, "palette_mode")
        column.prop(settings, "apply_modifiers")

        export_copy = layout.box()
        export_copy.prop(settings, "create_export_copy")
        quick_export = export_copy.column(align=True)
        quick_export.enabled = settings.create_export_copy
        quick_export.prop(settings, "quick_export_glb")
        if settings.quick_export_glb:
            quick_export.prop(settings, "glb_filepath")

        output = layout.box()
        output.label(text="Output")
        output.prop(settings, "image_name")
        output.prop(settings, "save_to_disk")
        if settings.save_to_disk:
            output.prop(settings, "filepath")
            output.prop(settings, "write_legend")

        operator_column = layout.column()
        operator_column.enabled = bool(mesh.uv_layers and context.mode == "OBJECT")
        operator_column.operator("object.bake_material_id", icon="RENDER_STILL")
        if context.mode != "OBJECT":
            layout.label(text="Switch to Object Mode to bake", icon="INFO")

        if settings.last_palette:
            palette_box = layout.box()
            palette_box.label(text=f"Last Bake: {settings.last_image_name}")
            for entry in settings.last_palette:
                row = palette_box.row(align=True)
                row.label(text=f"{entry.slot}: {entry.material_name}")
                row.prop(entry, "color", text="")


_CLASSES = (
    MIDB_PaletteEntry,
    MIDB_Settings,
    MIDB_OT_Bake,
    MIDB_PT_Panel,
)


def register() -> None:
    for cls in _CLASSES:
        bpy.utils.register_class(cls)
    bpy.types.Scene.material_id_baker = PointerProperty(type=MIDB_Settings)


def unregister() -> None:
    del bpy.types.Scene.material_id_baker
    for cls in reversed(_CLASSES):
        bpy.utils.unregister_class(cls)
