import bpy
import os
from mathutils import Vector, Matrix
import math
import sys

# --- Configuration ---
CONFIG = {
    "output_directory": "C:/Users/Joshua/Documents/GitHub/BA/LPTK_ASSETPACK/Thumbnails",
    "thumbnail_resolution_x": 256,
    "thumbnail_resolution_y": 256,
    "camera_distance_factor": 1.4,
    "camera_angle_x": math.radians(30), # Camera rotation around X-axis (up/down)
    "camera_angle_z": math.radians(25), # Camera rotation around Z-axis (left/right)
    "light_power": 8
}

# --- Helper Functions ---

def cleanup_scene_elements(names_and_collections: list):
    """
    Removes specified objects and their associated data blocks from the scene.
    Args:
        names_and_collections: A list of tuples, where each tuple is (object_name, data_collection).
                               data_collection should be bpy.data.cameras, bpy.data.lights, etc.
    """
    for obj_name, data_collection in names_and_collections:
        # 1. Remove Object
        obj_to_remove = bpy.data.objects.get(obj_name)
        if obj_to_remove:
            # Unlink from all collections before removing
            for collection in list(obj_to_remove.users_collection):
                if obj_to_remove.name in collection.objects:
                    collection.objects.unlink(obj_to_remove)
            bpy.data.objects.remove(obj_to_remove, do_unlink=True)

        # 2. Remove Data Block (if no other users)
        data_to_remove = data_collection.get(obj_name)
        if data_to_remove and data_to_remove.users == 0:
            data_collection.remove(data_to_remove)
        elif data_to_remove and data_to_remove.users > 0:
            print(f"Warning: Data block '{obj_name}' still has {data_to_remove.users} users and was not removed.")

def setup_camera_and_light():
    """
    Sets up a new camera and sun lamp for thumbnail rendering.
    Returns:
        (bpy.types.Object, bpy.types.Object): The camera object and sun object.
    """
    # Camera
    cam_data = bpy.data.cameras.new(name='ThumbnailCam')
    cam_obj = bpy.data.objects.new('ThumbnailCam', cam_data)
    bpy.context.scene.collection.objects.link(cam_obj)
    bpy.context.scene.camera = cam_obj

    cam_obj.data.type = 'PERSP'
    cam_obj.data.lens = 50

    # Sun Light
    sun_data = bpy.data.lights.new(name='ThumbnailSun', type='SUN')
    sun_data.energy = CONFIG["light_power"]
    sun_data.angle = math.radians(10) # Soften shadows

    sun_obj = bpy.data.objects.new('ThumbnailSun', sun_data)
    bpy.context.scene.collection.objects.link(sun_obj)
    sun_obj.location = (5, -5, 5) # General position
    sun_obj.rotation_euler = (math.radians(45), math.radians(-30), math.radians(15)) # Angle it

    print("Camera and light set up.")
    return cam_obj, sun_obj

def position_camera_around_object(obj, cam_obj):
    """
    Positions the camera to frame the given object based on its bounding box.
    """
    cam = cam_obj.data

    if not obj.bound_box:
        print(f"Skipping camera positioning for {obj.name}: Mesh has no geometry contributing to bounding box.")
        return False # Indicate failure to position

    bbox_corners_world = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box]
    bbox_center = sum(bbox_corners_world, Vector()) / 8.0

    min_x = min(v.x for v in bbox_corners_world)
    max_x = max(v.x for v in bbox_corners_world)
    min_y = min(v.y for v in bbox_corners_world)
    max_y = max(v.y for v in bbox_corners_world)
    min_z = min(v.z for v in bbox_corners_world)
    max_z = max(v.z for v in bbox_corners_world)

    width = max_x - min_x
    height = max_z - min_z
    depth = max_y - min_y

    max_dim = max(width, height, depth)
    if max_dim == 0: # Handle cases where all dimensions are zero (e.g., single vertex)
        max_dim = 0.01 # Provide a tiny default size

    # Calculate distance based on camera FoV
    if cam.angle == 0:
        # Prevent division by zero if camera angle is exactly 0; use a tiny angle instead
        distance = max_dim * CONFIG["camera_distance_factor"] / (2 * math.tan(math.radians(0.001) / 2))
    else:
        distance = max_dim * CONFIG["camera_distance_factor"] / (2 * math.tan(cam.angle / 2))

    # Calculate camera position based on angles
    cam_x = distance * math.sin(CONFIG["camera_angle_z"]) * math.cos(CONFIG["camera_angle_x"])
    cam_y = -distance * math.cos(CONFIG["camera_angle_z"]) * math.cos(CONFIG["camera_angle_x"]) # Negative Y for front view
    cam_z = distance * math.sin(CONFIG["camera_angle_x"])

    cam_obj.location = bbox_center + Vector((cam_x, cam_y, cam_z))

    # Aim camera at the bounding box center
    direction = bbox_center - cam_obj.location
    if direction.length_squared == 0:
        cam_obj.rotation_euler = (0,0,0) # Default rotation if no clear direction
    else:
        rot_quat = direction.to_track_quat('-Z', 'Y') # Point -Z axis towards target, Y as up
        cam_obj.rotation_euler = rot_quat.to_euler()

    bpy.context.view_layer.update()
    return True

def render_and_assign_thumbnail(obj, output_dir):
    """
    Renders a thumbnail for the given object and attempts to assign it as its asset preview
    using bpy.ops.ed.lib_id_load_custom_preview().
    """
    scene = bpy.context.scene
    
    filepath = os.path.join(output_dir, f"{obj.name}Thumbnail.png")
    
    scene.render.filepath = filepath
    bpy.ops.render.render(write_still=True)
    print(f"Rendered thumbnail for: {obj.name} -> {os.path.basename(filepath)}")

    if obj.asset_data:
        # Use the operator directly with temp_override
        print(f"\n--- ATTEMPTING TO LOAD CUSTOM PREVIEW FOR '{obj.name}' USING OPERATOR ---")
        try:
            # The 'id' argument to temp_override should be the ID block for which the preview is being loaded.
            # For objects, the object itself (bpy.types.Object) is the ID block holding the asset_data.
            id_block = obj

            # Make sure the file exists before calling the operator
            if not os.path.exists(filepath):
                print(f"ERROR: Thumbnail file '{filepath}' does not exist. Cannot load custom preview for {obj.name}.")
                return

            # Override context to ensure the operator acts on the correct ID block
            with bpy.context.temp_override(id=id_block):
                bpy.ops.ed.lib_id_load_custom_preview(filepath=filepath)
            
            # After successful operation, check if the image is now linked and packed
            if hasattr(obj.asset_data, 'preview') and obj.asset_data.preview:
                print(f"SUCCESS: Custom preview loaded and assigned for {obj.name} using operator.")
                if obj.asset_data.preview.packed_file:
                    print(f"Image '{obj.asset_data.preview.name}' is packed into .blend file.")
                else:
                    # If the operator didn't pack it (unlikely for this specific op), we could try manually.
                    # However, usually this operator takes care of packing.
                    print(f"WARNING: Image '{obj.asset_data.preview.name}' was NOT packed by operator. This might indicate an issue with packing.")
            else:
                print(f"WARNING: Operator ran, but 'preview' attribute still not found or assigned for {obj.name}.")
                print(f"This indicates a deeper issue where even the operator cannot set the preview correctly.")

        except Exception as e:
            print(f"ERROR: Failed to load custom preview for {obj.name} using operator: {e}")
            print(f"Could not assign custom preview for {obj.name}.")
    else:
        print(f"Warning: {obj.name} is not marked as an asset. Skipping preview assignment.")

# --- Main Script Execution ---
def main():
    # Ensure output directory exists
    if not os.path.exists(CONFIG["output_directory"]):
        os.makedirs(CONFIG["output_directory"])
        print(f"Created output directory: {CONFIG['output_directory']}")

    # Configure scene render settings
    scene = bpy.context.scene
    scene.render.engine = 'BLENDER_EEVEE_NEXT' # Using EEVEE as it's standard in 4.3
    scene.render.film_transparent = True
    scene.render.image_settings.file_format = 'PNG'
    scene.render.image_settings.color_mode = 'RGBA'
    scene.render.resolution_x = CONFIG["thumbnail_resolution_x"]
    scene.render.resolution_y = CONFIG["thumbnail_resolution_y"]
    scene.render.resolution_percentage = 100

    # Cleanup existing camera and light from previous runs
    cleanup_scene_elements([
        ("ThumbnailCam", bpy.data.cameras),
        ("ThumbnailSun", bpy.data.lights)
    ])

    # Setup new camera and light
    cam_obj, sun_obj = setup_camera_and_light()

    # Filter for actual mesh objects that are considered assets in the current file
    asset_objects = [o for o in bpy.data.objects if o.asset_data and o.type == 'MESH']

    if not asset_objects:
        print("\nNo mesh objects with asset data found in the scene. Please mark your objects as assets.")
        # Restore visibility even if no assets found
        for obj in bpy.data.objects:
            obj.hide_render = False
        return # Exit the function if nothing to process

    print(f"\nProcessing {len(asset_objects)} asset(s)...")

    # Process each asset object
    for obj in asset_objects:
        print(f"\n--- Processing Asset: {obj.name} ---")
        bpy.ops.object.select_all(action='DESELECT') # Deselect all
        # We don't need to select/activate for temp_override(id=obj) for lib_id_load_custom_preview
        # obj.select_set(True)
        # bpy.context.view_layer.objects.active = obj

        # Hide all other renderable objects except the current asset, camera, and lights
        for other_obj in bpy.data.objects:
            if other_obj != obj and other_obj != cam_obj and other_obj != sun_obj:
                other_obj.hide_render = True
        obj.hide_render = False # Ensure the asset itself is renderable

        # Position camera and render
        if obj.type == 'MESH' and obj.data and obj.data.vertices:
            if position_camera_around_object(obj, cam_obj):
                render_and_assign_thumbnail(obj, CONFIG["output_directory"])
            else:
                print(f"Skipping thumbnail generation and assignment for {obj.name} due to camera positioning issue.")
        else:
            print(f"Skipping {obj.name}: Not a valid mesh type or has no vertex data to process.")
        
        # Restore visibility of the current object immediately after processing
        obj.hide_render = False 

    # Final Cleanup / Restore Visibility of all objects (just in case any were missed)
    print("\nRestoring visibility of all objects...")
    for obj in bpy.data.objects:
        obj.hide_render = False

    print("\n--- Script Finished ---")
    print(f"Thumbnails saved to: {CONFIG['output_directory']}")

# Run the main function
if __name__ == "__main__":
    main()