import bpy
from mathutils import Matrix
from typing import List, Optional
from dataclasses import dataclass

@dataclass
class Snapshot:
    m1: Optional[Matrix];
    m2: Optional[Matrix];
    f1: int;
    f2: int;
    obj: Optional[bpy.types.Object];

@dataclass
class TimeCache:
    obj: bpy.types.Object;
    mat: Matrix;
    cache: List[Snapshot];

cached_objects:List[TimeCache] = []
blocking = False

def sample_transform(obj:bpy.types.Object, frame: int|float)->Matrix:
    
    """Return the object's world transform matrix at a given frame number."""
    if obj is None:
        raise ValueError("No object provided")
    
    scene = bpy.context.scene
    
    if frame != scene.frame_current:
    
        global blocking
        blocking = True
        
        depsgraph = bpy.context.evaluated_depsgraph_get()

        current_frame = scene.frame_current
        scene.frame_set(int(frame))
        eval_obj = obj.evaluated_get(depsgraph)
        matrix = eval_obj.matrix_world.copy()
        scene.frame_set(current_frame)
        
        blocking = False
        
        return matrix
    
    else:
        return obj.matrix_world.copy()

def apply_accumulated_deltas(data:TimeCache, sample_time:int)->None:
    """
    obj_data = (obj, base_matrix, cached_segments)
    Applies deltas to obj based on sample_time.
    """
    # obj, base_matrix, cached_data = obj_data


    cumulative = Matrix.Identity(4)
    last_range_end = None

    for snapshot in data.cache:
        if snapshot.obj is None:
            continue
        if sample_time >= snapshot.f2:
            delta = snapshot.m2 @ snapshot.m1.inverted()
            cumulative = delta @ cumulative
            last_range_end = snapshot.f2
        elif snapshot.f1 <= sample_time < snapshot.f2:
            current = sample_transform(snapshot.obj, sample_time)
            delta = current @ snapshot.m1.inverted()
            cumulative = delta @ cumulative
            last_range_end = sample_time
            break
        else:
            break

    if last_range_end is None:
        return

    data.obj.matrix_world = cumulative @ data.mat

def parse_scene_objects() -> None:
    """
    Scans all objects in the current scene for Child Of constraints named like 'frame_<frame_number>'.
    Records for each object:
        - Base matrix captured one frame before the first control
        - List of control segments (m1, m2, f1, f2, control)
    """
    global cached_objects
    cached_objects = []
    for target in bpy.context.scene.objects:
        parse_object(target)
    
def parse_object(target:bpy.types.Object)->Optional[TimeCache]:
    
    frame_entries:List[Snapshot] = []
    
    for constraint in target.constraints:
        if constraint.type == 'CHILD_OF' and constraint.name.startswith("frame_"):
            try:
                frame = int(constraint.name.split("_")[1])
                frame_entries.append(Snapshot(m1=None, m2=None, f1=frame, f2=frame, obj=constraint.target))
            except Exception as e:
                print(f"Invalid constraint name '{constraint.name}' on '{target.name}': {e}")
            finally:
                constraint.mute = True  # disable constraint

    if len(frame_entries) == 0:
        return None

    frame_entries.sort(key=lambda x: x.f1)

    # capture baseline before first frame
    first_frame:int = frame_entries[0].f1
    base_frame:int = max(first_frame - 1, 0)
    base_matrix:Matrix = sample_transform(target, base_frame)

    #for i, (f1, control) in enumerate(frame_entries):
    for i, snapshot in enumerate(frame_entries):
        check_next = frame_entries[i + 1].f1 if i + 1 < len(frame_entries) else 999999
        snapshot.f2 = check_next - 1 if check_next > snapshot.f1 else snapshot.f1

        print("frames are:", snapshot.f1, snapshot.f2)

        if snapshot.obj:
            snapshot.m1 = sample_transform(snapshot.obj, snapshot.f1)
            snapshot.m2 = sample_transform(snapshot.obj, snapshot.f2)

    cached = TimeCache(obj=target, mat=base_matrix, cache=frame_entries)
    global cached_objects
    cached_objects.append(cached)
    return cached
    
def release_object(obj:bpy.types.Object)->None:

    global cached_objects

    for i, record in enumerate(cached_objects):
        if obj is record.obj:
            del cached_objects[i]
            return

def recompute()->None:
    for obj_data in cached_objects:
        apply_accumulated_deltas(obj_data, bpy.context.scene.frame_current)

def frame_change_handler(scene:bpy.types.Scene)->None:
    if blocking:
        print("blocked")
    else:
        recompute()

def unregister_existing_handler()->None:
    # Remove any previously registered version of this handler
    for h in bpy.app.handlers.frame_change_post:
        if h.__name__ == "frame_change_handler":
            bpy.app.handlers.frame_change_post.remove(h)
            print("Removed old frame_change_handler")

def register_handler()->None:
    unregister_existing_handler()
    bpy.app.handlers.frame_change_post.append(frame_change_handler)
    print("Registered new frame_change_handler")