"""
bake_export.py
--------------
Operator to bake the Pr3tz procedural animation into a fully
self-contained .blend file suitable for distributed render farms
(SheepIt, etc.).

The resulting file contains:
  - One mesh object per unique frame, with hide_render / hide_viewport
    keyframes so only the correct mesh is visible on each frame.
  - All materials (preset, blend, project) referenced by the objects.
  - Camera with baked delta_location keyframes (camera shake).
  - Light, world, and render settings copied from the source scene.
  - No addon dependencies — the file is fully standalone.
"""
from __future__ import annotations

import hashlib
import random
import struct

import bpy
from bpy_extras.io_utils import ExportHelper

from .constants import KNOT_OBJ_NAME


# ---------------------------------------------------------------------------
# Blender 5.0 — layered Action API (Action.fcurves removed)
# ---------------------------------------------------------------------------

def _get_action_fcurves(obj):
    """Return an iterable of FCurves for the given object's active action slot.

    Blender 5.0 uses a layered animation system:
      action → layers → strips → channelbag (per slot) → fcurves
    """
    anim = obj.animation_data
    if not anim or not anim.action:
        return []

    action = anim.action
    slot   = anim.action_slot
    if slot is None:
        return []

    for layer in action.layers:
        for strip in layer.strips:
            cb = strip.channelbag(slot)
            if cb is not None:
                return cb.fcurves
    return []



# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _mesh_fingerprint(mesh: bpy.types.Mesh) -> str:
    """Compute a fast hash of a mesh's vertex positions + material index.

    Used to detect consecutive identical frames so we can extend the
    previous object's visibility window instead of creating a duplicate.
    """
    h = hashlib.md5(usedforsecurity=False)

    # Hash vertex positions
    for v in mesh.vertices:
        h.update(struct.pack('fff', *v.co))

    # Hash which material slots are assigned
    for slot_idx, slot in enumerate(mesh.materials):
        name = slot.name if slot else ""
        h.update(f"{slot_idx}:{name}".encode())

    return h.hexdigest()


def _obj_transform_fingerprint(obj: bpy.types.Object) -> str:
    """Hash an object's location, rotation_euler, and scale."""
    h = hashlib.md5(usedforsecurity=False)
    h.update(struct.pack('fff', *obj.location))
    h.update(struct.pack('fff', *obj.rotation_euler))
    h.update(struct.pack('fff', *obj.scale))
    return h.hexdigest()


def _keyframe_visibility(obj, frame: int, visible: bool) -> None:
    """Insert hide_render and hide_viewport keyframes for a single frame."""
    obj.hide_render   = not visible
    obj.hide_viewport = not visible
    obj.keyframe_insert(data_path="hide_render",   frame=frame)
    obj.keyframe_insert(data_path="hide_viewport", frame=frame)


def _set_constant_interpolation(obj) -> None:
    """Set all visibility F-curves to CONSTANT interpolation.

    This prevents Blender from smoothly fading between True/False,
    which would make objects partially visible on intermediate frames.
    """
    for fc in _get_action_fcurves(obj):
        if fc.data_path in ("hide_render", "hide_viewport"):
            for kp in fc.keyframe_points:
                kp.interpolation = 'CONSTANT'


# ---------------------------------------------------------------------------
# Main bake operator
# ---------------------------------------------------------------------------

class KNOT_OT_BakeExport(bpy.types.Operator, ExportHelper):
    """Bake the procedural animation into a standalone .blend file."""
    bl_idname    = "knot.bake_export"
    bl_label     = "Bake & Export for Render Farm"
    bl_description = (
        "Iterate every frame, convert the procedural knot to mesh, "
        "and save a self-contained .blend with per-frame visibility keyframes. "
        "The output file requires no addons and is render-farm ready"
    )
    bl_options = {'REGISTER'}

    # ExportHelper provides the file browser
    filename_ext = ".blend"
    filter_glob: bpy.props.StringProperty(default="*.blend", options={'HIDDEN'})

    # User-facing options shown in the file browser sidebar
    use_render_resolution: bpy.props.BoolProperty(
        name="Use Render Resolution",
        description="Bake at the Render Curve / Bevel resolution (recommended for final output)",
        default=True,
    )
    pack_textures: bpy.props.BoolProperty(
        name="Pack Textures",
        description="Pack all external textures into the .blend file",
        default=True,
    )
    split_export: bpy.props.BoolProperty(
        name="Split Export",
        description="Split the baked animation into multiple .blend files",
        default=False,
    )
    frames_per_file: bpy.props.IntProperty(
        name="Frames Per File",
        description="Number of frames to include in each split file",
        default=500,
        min=1,
    )

    def execute(self, context):
        scene = context.scene

        # Validate
        if len(scene.knot_list) == 0:
            self.report({'ERROR'}, "Playlist is empty — nothing to bake.")
            return {'CANCELLED'}

        knot_obj = bpy.data.objects.get(KNOT_OBJ_NAME)
        if not knot_obj:
            self.report({'ERROR'}, f"AnimKnot object '{KNOT_OBJ_NAME}' not found.")
            return {'CANCELLED'}

        filepath = self.filepath
        if not filepath.lower().endswith(".blend"):
            filepath += ".blend"

        orig_frame_start = scene.frame_start
        orig_frame_end   = scene.frame_end
        total_frames = orig_frame_end - orig_frame_start + 1

        self.report({'INFO'}, f"Baking {total_frames} frames...")

        chunk_size = self.frames_per_file if self.split_export else total_frames

        # We will loop over chunks and bake each as a separate file
        try:
            import os
            first_chunk = True
            for chunk_start in range(orig_frame_start, orig_frame_end + 1, chunk_size):
                chunk_end = min(chunk_start + chunk_size - 1, orig_frame_end)

                scene.frame_start = chunk_start
                scene.frame_end = chunk_end

                chunk_filepath = filepath
                if self.split_export:
                    base, ext = os.path.splitext(filepath)
                    chunk_filepath = f"{base}_{chunk_start}-{chunk_end}{ext}"

                self.report({'INFO'}, f"Baking chunk {chunk_start}-{chunk_end} to {chunk_filepath}...")
                # pack_and_clean is True only on the first chunk so that external
                # textures are packed exactly once.  Subsequent chunks already
                # have the packed images in memory from the first run.
                result = self._bake_range(
                    context, chunk_filepath,
                    pack_and_clean=first_chunk,
                )
                first_chunk = False
                if result != {'FINISHED'}:
                    self.report({'ERROR'}, f"Failed baking chunk {chunk_start}-{chunk_end}")
                    return {'CANCELLED'}
        finally:
            scene.frame_start = orig_frame_start
            scene.frame_end = orig_frame_end

        self.report({'INFO'}, "Baking complete.")
        return {'FINISHED'}

    def _bake_range(self, context, filepath, *, pack_and_clean: bool = True):
        scene = context.scene
        glob  = scene.knot_globals

        frame_start = scene.frame_start
        frame_end   = scene.frame_end
        total_frames = frame_end - frame_start + 1

        # If "Use Render Resolution" is on, temporarily override the preview
        # resolution to match render resolution.  The handler reads
        # preview_resolution for viewport and render_resolution for renders,
        # but frame_set() triggers a viewport-mode handler call.  By
        # equalising them we ensure the baked mesh uses render quality.
        orig_preview_res   = glob.preview_resolution
        orig_preview_bevel = glob.preview_bevel_resolution
        if self.use_render_resolution:
            glob.preview_resolution       = glob.render_resolution
            glob.preview_bevel_resolution = glob.render_bevel_resolution

        # -------------------------------------------------------------------
        # Phase 1: Bake per-frame geometry into mesh objects
        # -------------------------------------------------------------------

        # Collect all baked objects so we can keyframe them afterwards
        baked_objects: list[tuple[bpy.types.Object, int, int]] = []
        # Each entry: (object, first_visible_frame, last_visible_frame)

        prev_fingerprint = None
        prev_transform_fp = None

        # Create a dedicated collection to hold baked frames
        bake_coll = bpy.data.collections.new("BakedAnimation")
        scene.collection.children.link(bake_coll)

        # Store the original frame so we can restore it later
        original_frame = scene.frame_current

        for f in range(frame_start, frame_end + 1):
            # Set the frame — this triggers knot_frame_handler via
            # the registered frame_change_post handler
            scene.frame_set(f)

            # Get the knot object (may have been recreated by the handler)
            knot_obj = bpy.data.objects.get(KNOT_OBJ_NAME)
            if not knot_obj:
                continue

            # Evaluate the depsgraph to get the final geometry
            depsgraph = context.evaluated_depsgraph_get()
            eval_obj  = knot_obj.evaluated_get(depsgraph)

            # Convert to mesh
            mesh = bpy.data.meshes.new_from_object(eval_obj, depsgraph=depsgraph)
            mesh.name = f"BakedKnot_Mesh_F{f}"

            # Compute fingerprints for duplicate detection
            mesh_fp      = _mesh_fingerprint(mesh)
            transform_fp = _obj_transform_fingerprint(knot_obj)

            if (mesh_fp == prev_fingerprint
                    and transform_fp == prev_transform_fp
                    and baked_objects):
                # Identical to the previous frame — extend the previous
                # object's visibility window and discard the duplicate mesh
                bpy.data.meshes.remove(mesh)
                _, first_f, _ = baked_objects[-1]
                baked_objects[-1] = (baked_objects[-1][0], first_f, f)
                continue

            prev_fingerprint  = mesh_fp
            prev_transform_fp = transform_fp

            # Create a new object for this frame
            obj = bpy.data.objects.new(f"BakedKnot_F{f}", mesh)

            # Copy transform
            obj.location       = knot_obj.location.copy()
            obj.rotation_euler = knot_obj.rotation_euler.copy()
            obj.scale          = knot_obj.scale.copy()

            # Copy material slots
            for src_slot in knot_obj.material_slots:
                obj.data.materials.append(src_slot.material)

            # Copy display settings
            obj.display_type = knot_obj.display_type

            bake_coll.objects.link(obj)
            baked_objects.append((obj, f, f))

        # Restore preview resolution if we overrode it
        if self.use_render_resolution:
            glob.preview_resolution       = orig_preview_res
            glob.preview_bevel_resolution = orig_preview_bevel

        # Restore original frame
        scene.frame_set(original_frame)

        if not baked_objects:
            self.report({'ERROR'}, "No frames were baked — the handler may not be producing geometry.")
            # Clean up the empty collection
            bpy.data.collections.remove(bake_coll)
            return {'CANCELLED'}

        # -------------------------------------------------------------------
        # Phase 2: Set up visibility keyframes
        # -------------------------------------------------------------------

        for obj, first_f, last_f in baked_objects:
            # Default: hidden on all frames
            # We keyframe "hidden" before and after the active window,
            # and "visible" for the active range.

            # Hidden before the active range
            if first_f > frame_start:
                _keyframe_visibility(obj, frame_start, False)
            if first_f > frame_start + 1:
                _keyframe_visibility(obj, first_f - 1, False)

            # Visible during the active range
            _keyframe_visibility(obj, first_f, True)
            if last_f != first_f:
                _keyframe_visibility(obj, last_f, True)

            # Hidden after the active range
            if last_f < frame_end:
                _keyframe_visibility(obj, last_f + 1, False)
            if last_f < frame_end - 1:
                _keyframe_visibility(obj, frame_end, False)

            _set_constant_interpolation(obj)

        # -------------------------------------------------------------------
        # Phase 3: Bake camera shake to keyframes
        # -------------------------------------------------------------------

        cam = scene.camera
        if cam and glob.camera_shake_amplitude > 0.0:
            shake = glob.camera_shake_amplitude
            for f in range(frame_start, frame_end + 1):
                import random
                random.seed(int(f * 1000))
                cam.delta_location = (
                    random.uniform(-shake, shake),
                    random.uniform(-shake, shake),
                    random.uniform(-shake, shake),
                )
                cam.keyframe_insert(data_path="delta_location", frame=f)

            # Set interpolation to LINEAR for smooth shake
            for fc in _get_action_fcurves(cam):
                if fc.data_path == "delta_location":
                    for kp in fc.keyframe_points:
                        kp.interpolation = 'LINEAR'

        # If no shake, ensure delta_location is zeroed
        elif cam:
            cam.delta_location = (0.0, 0.0, 0.0)

        # -------------------------------------------------------------------
        # Phase 4: Clean up the procedural AnimKnot object
        # -------------------------------------------------------------------

        # Remove the original NURBS curve object — we've replaced it with
        # baked mesh objects
        knot_obj = bpy.data.objects.get(KNOT_OBJ_NAME)
        if knot_obj:
            curve_data = knot_obj.data
            bpy.data.objects.remove(knot_obj, do_unlink=True)
            if curve_data and curve_data.users == 0:
                bpy.data.curves.remove(curve_data)

        # Remove the frame_change_post handler so the baked file is
        # completely standalone
        from .handler import knot_frame_handler
        handlers = bpy.app.handlers.frame_change_post
        for h in list(handlers):
            if getattr(h, "__name__", "") == knot_frame_handler.__name__:
                handlers.remove(h)

        # -------------------------------------------------------------------
        # Phase 5: Pack data and ensure render-farm compatibility
        # -------------------------------------------------------------------

        # Set output format to PNG image sequence (required by SheepIt)
        orig_file_format = scene.render.image_settings.file_format
        orig_color_mode  = scene.render.image_settings.color_mode
        scene.render.image_settings.file_format = 'PNG'
        scene.render.image_settings.color_mode  = 'RGBA'

        if pack_and_clean:
            # Make all paths relative — only needed once; subsequent chunks
            # operate on the same in-memory data which is already relative.
            try:
                bpy.ops.file.make_paths_relative()
            except RuntimeError:
                pass  # May fail if file has never been saved

            # Pack external data if requested — run once so that all packed
            # images are available in memory for every subsequent chunk save.
            if self.pack_textures:
                try:
                    bpy.ops.file.pack_all()
                except RuntimeError:
                    self.report({'WARNING'}, "Some external files could not be packed.")

            # Remove unused data blocks to reduce file size.  Only safe to
            # do on the first chunk; later chunks share the same in-memory
            # materials/images and we do not want to inadvertently purge data
            # that is still referenced by a subsequent chunk's bake.
            for mat in list(bpy.data.materials):
                if mat.users == 0:
                    bpy.data.materials.remove(mat)
            for m in list(bpy.data.meshes):
                if m.users == 0:
                    bpy.data.meshes.remove(m)
            for c in list(bpy.data.curves):
                if c.users == 0:
                    bpy.data.curves.remove(c)
            for ng in list(bpy.data.node_groups):
                if ng.users == 0:
                    bpy.data.node_groups.remove(ng)
            for img in list(bpy.data.images):
                if img.users == 0 and not img.use_fake_user:
                    bpy.data.images.remove(img)

        # -------------------------------------------------------------------
        # Phase 6: Save as a COPY — the original session is untouched
        # -------------------------------------------------------------------

        bpy.ops.wm.save_as_mainfile(filepath=filepath, copy=True, compress=True)

        # -------------------------------------------------------------------
        # Phase 7: Restore the original session state
        # -------------------------------------------------------------------
        # We just saved a copy, but the in-memory state has been modified
        # (baked objects added, AnimKnot removed, handler removed).
        # Restore everything so the user can keep working.

        # Restore render settings
        scene.render.image_settings.file_format = orig_file_format
        scene.render.image_settings.color_mode  = orig_color_mode

        # Remove all baked objects and their meshes from the current session
        for obj, _, _ in baked_objects:
            if obj and obj.name in bpy.data.objects:
                mesh_data = obj.data
                bpy.data.objects.remove(obj, do_unlink=True)
                if mesh_data and mesh_data.users == 0:
                    bpy.data.meshes.remove(mesh_data)

        # Remove the baked collection
        bake_coll_check = bpy.data.collections.get("BakedAnimation")
        if bake_coll_check:
            bpy.data.collections.remove(bake_coll_check)

        # Remove baked camera shake keyframes (restore static camera)
        if cam:
            cam_fcurves = _get_action_fcurves(cam)
            # cam_fcurves may be a ChannelBag or ActionFCurves — both support .remove()
            fcurves_to_remove = [
                fc for fc in cam_fcurves
                if fc.data_path == "delta_location"
            ]
            for fc in fcurves_to_remove:
                try:
                    cam_fcurves.remove(fc)
                except (TypeError, RuntimeError):
                    pass  # Some API versions don't support removal — keyframes will persist
            cam.delta_location = (0.0, 0.0, 0.0)

        # Re-create the AnimKnot object and re-register the handler
        from .geometry import _make_torus_knot
        from .materials import prewarm_materials_and_blends
        
        # We need to prewarm materials and blends because removing unused materials might have deleted them
        prewarm_materials_and_blends(scene)

        # Re-register the handler
        handlers = bpy.app.handlers.frame_change_post
        already_registered = any(
            getattr(h, "__name__", "") == knot_frame_handler.__name__
            for h in handlers
        )
        if not already_registered:
            handlers.append(knot_frame_handler)

        # Trigger the handler once to rebuild the AnimKnot object
        knot_frame_handler(scene)

        # Restore the frame
        scene.frame_set(original_frame)

        return {'FINISHED'}

    def draw(self, context):
        """Draw options in the file browser sidebar."""
        layout = self.layout
        layout.label(text="Bake Options:", icon='SETTINGS')
        layout.prop(self, "use_render_resolution")
        layout.prop(self, "pack_textures")
        
        layout.separator()
        layout.prop(self, "split_export")
        if self.split_export:
            layout.prop(self, "frames_per_file")

        scene = context.scene
        total = scene.frame_end - scene.frame_start + 1
        layout.separator()
        layout.label(text=f"Frames: {scene.frame_start} – {scene.frame_end} ({total} total)")
        layout.label(text=f"Playlist: {len(scene.knot_list)} knots")

