import { BufferGeometry, Camera, Color, Euler,Group, Material, Object3D, Scene, Texture, Vector2, Vector3, Vector4, WebGLRenderer } from "three";


// @TODO: we need to detect objects with materials both transparent and NOT transparent. These need to be updated in scene.onBeforeRender to have correct renderlists

/**
 * Valid types that can be used as material property overrides
 */
type MaterialPropertyType = number | number[] | Color | Texture | Vector2 | Vector3 | Vector4 | null | Euler;

/**
 * Defines offset and repeat transformations for texture coordinates
 */
export interface TextureTransform {
    /** UV offset applied to the texture */
    offset?: Vector2;
    /** UV repeat/scale applied to the texture */
    repeat?: Vector2;
}

/**
 * Represents a single material property override with optional texture transformation
 * @template T The type of the property value
 */
export interface PropertyBlockOverride<T extends MaterialPropertyType = MaterialPropertyType> {
    /** The name of the material property to override (e.g., "color", "map", "roughness") */
    name: string;
    /** The value to set for this property */
    value: T;
    /** Optional texture coordinate transformation (only used when value is a Texture) */
    textureTransform?: TextureTransform;
}

/**
 * Utility type that extracts only non-function property names from a type
 * @template T The type to extract property names from
 */
type NonFunctionPropertyNames<T> = {
    [K in keyof T]: T[K] extends Function ? never : K
}[keyof T];

/**
 * Centralized registry for all property block related data.
 * Uses WeakMaps to allow automatic garbage collection when objects are destroyed.
 * @internal
 */
class PropertyBlockRegistry {
    // Map from object to its property block
    private objectToBlock = new WeakMap<Object3D, MaterialPropertyBlock>();

    // Track which materials belong to which property block (to prevent applying to wrong materials)
    // Use WeakSet for automatic cleanup when materials are garbage collected
    // private objectToMaterials = new WeakMap<Object3D, WeakSet<Material>>();

    // Track which meshes have callbacks for which property block owners
    private meshToOwners = new WeakMap<Object3D, Set<Object3D>>();

    // Track original callback functions for cleanup (reserved for future use)
    private meshToOriginalCallbacks = new WeakMap<Object3D, {
        onBeforeRender?: ObjectRenderCallback;
        onAfterRender?: ObjectRenderCallback;
    }>();

    getBlock(object: Object3D): MaterialPropertyBlock | undefined {
        return this.objectToBlock.get(object);
    }

    setBlock(object: Object3D, block: MaterialPropertyBlock): void {
        this.objectToBlock.set(object, block);
    }

    deleteBlock(object: Object3D): void {
        this.objectToBlock.delete(object);
        // this.objectToMaterials.delete(object);
    }

    // addMaterial(object: Object3D, material: Material): void {
    //     let materials = this.objectToMaterials.get(object);
    //     if (!materials) {
    //         materials = new WeakSet();
    //         this.objectToMaterials.set(object, materials);
    //     }
    //     materials.add(material);
    // }

    // hasMaterial(object: Object3D, material: Material): boolean {
    //     return this.objectToMaterials.get(object)?.has(material) ?? false;
    // }

    isHooked(mesh: Object3D, owner: Object3D): boolean {
        return this.meshToOwners.get(mesh)?.has(owner) ?? false;
    }

    addHook(mesh: Object3D, owner: Object3D): void {
        let owners = this.meshToOwners.get(mesh);
        if (!owners) {
            owners = new Set();
            this.meshToOwners.set(mesh, owners);
        }
        owners.add(owner);
    }

    removeHook(mesh: Object3D, owner: Object3D): void {
        const owners = this.meshToOwners.get(mesh);
        if (owners) {
            owners.delete(owner);
            if (owners.size === 0) {
                this.meshToOwners.delete(mesh);
            }
        }
    }

    getOriginalCallbacks(mesh: Object3D): { onBeforeRender?: ObjectRenderCallback; onAfterRender?: ObjectRenderCallback } | undefined {
        return this.meshToOriginalCallbacks.get(mesh);
    }

    setOriginalCallbacks(mesh: Object3D, callbacks: { onBeforeRender?: ObjectRenderCallback; onAfterRender?: ObjectRenderCallback }): void {
        this.meshToOriginalCallbacks.set(mesh, callbacks);
    }
}

const registry = new PropertyBlockRegistry();

/**
 * MaterialPropertyBlock allows per-object material property overrides without creating new material instances.
 * This is useful for rendering multiple objects with the same base material but different properties
 * (e.g., different colors, textures, or shader parameters).
 *
 * ## How Property Blocks Work
 *
 * **Important**: Overrides are registered on the **Object3D**, not on the material.
 * This means:
 * - If you change the object's material, the overrides will still be applied to the new material
 * - Multiple objects can share the same material but have different property overrides
 * - If you don't want overrides applied after changing a material, you must remove them using {@link removeOveride}, {@link clearAllOverrides}, or {@link dispose}
 *
 * The property block system works by:
 * - Temporarily applying overrides in onBeforeRender
 * - Restoring original values in onAfterRender
 * - Managing shader defines and program cache keys for correct shader compilation
 * - Supporting texture coordinate transforms per object
 *
 * ## Common Use Cases
 *
 * - **Lightmaps**: Apply unique lightmap textures to individual objects sharing the same material
 * - **Reflection Probes**: Apply different environment maps per object for localized reflections
 * - **See-through effects**: Temporarily override transparency/transmission properties for X-ray effects
 *
 * ## Getting a MaterialPropertyBlock
 *
 * **Important**: Do not use the constructor directly. Instead, use the static {@link MaterialPropertyBlock.get} method:
 *
 * ```typescript
 * const block = MaterialPropertyBlock.get(myMesh);
 * ```
 *
 * This method will either return an existing property block or create a new one if it doesn't exist.
 * It automatically:
 * - Creates the property block instance
 * - Registers it in the internal registry
 * - Attaches the necessary render callbacks to the object
 * - Handles Groups by applying overrides to all child meshes
 *
 * @example Basic usage
 * ```typescript
 * // Get or create a property block for an object
 * const block = MaterialPropertyBlock.get(myMesh);
 *
 * // Override the color property
 * block.setOverride("color", new Color(1, 0, 0));
 *
 * // Override a texture with custom UV transform (useful for lightmaps)
 * block.setOverride("lightMap", myLightmapTexture, {
 *   offset: new Vector2(0.5, 0.5),
 *   repeat: new Vector2(2, 2)
 * });
 *
 * // Set a shader define
 * block.setDefine("USE_CUSTOM_FEATURE", 1);
 * ```
 *
 * @example Material swapping behavior
 * ```typescript
 * const mesh = new Mesh(geometry, materialA);
 * const block = MaterialPropertyBlock.get(mesh);
 * block.setOverride("color", new Color(1, 0, 0));
 *
 * // The color override is red for materialA
 *
 * // Swap the material - overrides persist and apply to the new material!
 * mesh.material = materialB;
 * // The color override is now red for materialB too
 *
 * // If you don't want overrides on the new material, remove them:
 * block.clearAllOverrides(); // Remove all overrides
 * // or
 * block.removeOveride("color"); // Remove specific override
 * // or
 * block.dispose(); // Remove the entire property block
 * ```
 *
 * @example Lightmap usage
 * ```typescript
 * const block = MaterialPropertyBlock.get(mesh);
 * block.setOverride("lightMap", lightmapTexture);
 * block.setOverride("lightMapIntensity", 1.5);
 * ```
 *
 * @example See-through effect
 * ```typescript
 * const block = MaterialPropertyBlock.get(mesh);
 * block.setOverride("transparent", true);
 * block.setOverride("opacity", 0.3);
 * ```
 *
 * @template T The material type this property block is associated with
 */
export class MaterialPropertyBlock<T extends Material = Material> {
    private _overrides: PropertyBlockOverride[] = [];
    private _defines: Record<string, string | number | boolean> = {};
    private _object: Object3D | null = null;

    /** The object this property block is attached to */
    get object(): Object3D | null { return this._object; }

    /**
     * Creates a new MaterialPropertyBlock
     * @param object The object this property block is for (optional)
     */
    protected constructor(object: Object3D | null = null) {
        this._object = object;
    }

    /**
     * Gets or creates a MaterialPropertyBlock for the given object.
     * This is the recommended way to obtain a property block instance.
     *
     * @template T The material type
     * @param object The object to get/create a property block for
     * @returns The MaterialPropertyBlock associated with this object
     *
     * @example
     * ```typescript
     * const block = MaterialPropertyBlock.get(myMesh);
     * block.setOverride("roughness", 0.5);
     * ```
     */
    static get<T extends Material = Material>(object: Object3D): MaterialPropertyBlock<T> {
        let block = registry.getBlock(object);
        if (!block) {
            block = new MaterialPropertyBlock(object);
            registry.setBlock(object, block);
            attachPropertyBlockToObject(object, block);
        }
        return block as MaterialPropertyBlock<T>;
    }

    /**
     * Checks if an object has any property overrides
     * @param object The object to check
     * @returns True if the object has a property block with overrides
     */
    static hasOverrides(object: Object3D): boolean {
        const block = registry.getBlock(object);
        return block ? block.hasOverrides() : false;
    }

    /**
     * Disposes this property block and cleans up associated resources.
     * After calling dispose, this property block should not be used.
     */
    dispose() {
        if (this._object) {
            registry.deleteBlock(this._object);
            // TODO: Add cleanup for hooked meshes
        }
        this._overrides = [];
        this._object = null;
    }

    /**
     * Sets or updates a material property override.
     * The override will be applied to the material during rendering.
     *
     * @param name The name of the material property to override (e.g., "color", "map", "roughness")
     * @param value The value to set
     * @param textureTransform Optional UV transform (only used when value is a Texture)
     *
     * @example
     * ```typescript
     * // Override a simple property
     * block.setOverride("roughness", 0.8);
     *
     * // Override a color
     * block.setOverride("color", new Color(0xff0000));
     *
     * // Override a texture with UV transform
     * block.setOverride("map", texture, {
     *   offset: new Vector2(0, 0),
     *   repeat: new Vector2(2, 2)
     * });
     * ```
     */
    setOverride<K extends NonFunctionPropertyNames<T>>(name: K, value: T[K], textureTransform?: TextureTransform): void;
    setOverride(name: string, value: MaterialPropertyType, textureTransform?: TextureTransform): void;
    setOverride(name: string, value: MaterialPropertyType, textureTransform?: TextureTransform): void {
        const existing = this._overrides.find(o => o.name === name);
        if (existing) {
            existing.value = value;
            existing.textureTransform = textureTransform;
        } else {
            this._overrides.push({ name, value, textureTransform });
        }
    }

    /**
     * Gets the override for a specific property with type-safe value inference
     * @param name The property name to get
     * @returns The PropertyBlockOverride with correctly typed value if it exists, undefined otherwise
     *
     * @example
     * ```typescript
     * const block = MaterialPropertyBlock.get<MeshStandardMaterial>(mesh);
     *
     * // Value is inferred as number | undefined
     * const roughness = block.getOverride("roughness")?.value;
     *
     * // Value is inferred as Color | undefined
     * const color = block.getOverride("color")?.value;
     *
     * // Value is inferred as Texture | null | undefined
     * const map = block.getOverride("map")?.value;
     *
     * // Explicitly specify the type for properties not on the base material type
     * const transmission = block.getOverride<number>("transmission")?.value;
     *
     * // Or use a more specific material type
     * const physicalBlock = block as MaterialPropertyBlock<MeshPhysicalMaterial>;
     * const transmissionTyped = physicalBlock.getOverride("transmission")?.value; // number
     * ```
     */
    getOverride<K extends NonFunctionPropertyNames<T>>(name: K): PropertyBlockOverride<T[K] & MaterialPropertyType> | undefined;
    getOverride<V extends MaterialPropertyType = MaterialPropertyType>(name: string): PropertyBlockOverride<V> | undefined;
    getOverride(name: string): PropertyBlockOverride | undefined {
        return this._overrides.find(o => o.name === name);
    }

    /**
     * Removes a specific property override.
     * After removal, the material will use its original property value for this property.
     *
     * @param name The property name to remove the override for
     *
     * @example
     * ```typescript
     * const block = MaterialPropertyBlock.get(mesh);
     *
     * // Set some overrides
     * block.setOverride("color", new Color(1, 0, 0));
     * block.setOverride("roughness", 0.5);
     * block.setOverride("lightMap", lightmapTexture);
     *
     * // Remove a specific override - the material will now use its original color
     * block.removeOveride("color");
     *
     * // Other overrides (roughness, lightMap) remain active
     * ```
     */
    removeOveride<K extends NonFunctionPropertyNames<T>>(name: K | ({} & string)): void {
        const index = this._overrides.findIndex(o => o.name === name);
        if (index >= 0) {
            this._overrides.splice(index, 1);
        }
    }

    /**
     * Removes all property overrides from this block.
     * After calling this, the material will use its original values for all properties.
     *
     * **Note**: This does NOT remove shader defines. Use {@link clearDefine} or {@link dispose} for that.
     *
     * @example Remove all overrides but keep the property block
     * ```typescript
     * const block = MaterialPropertyBlock.get(mesh);
     *
     * // Set multiple overrides
     * block.setOverride("color", new Color(1, 0, 0));
     * block.setOverride("roughness", 0.5);
     * block.setOverride("lightMap", lightmapTexture);
     *
     * // Later, remove all overrides at once
     * block.clearAllOverrides();
     *
     * // The material now uses its original values
     * // The property block still exists and can be reused with new overrides
     * ```
     *
     * @example Temporarily disable all overrides
     * ```typescript
     * const block = MaterialPropertyBlock.get(mesh);
     *
     * // Save current overrides if you want to restore them later
     * const savedOverrides = [...block.overrides];
     *
     * // Clear all overrides temporarily
     * block.clearAllOverrides();
     *
     * // Do some rendering without overrides...
     *
     * // Restore overrides
     * savedOverrides.forEach(override => {
     *   block.setOverride(override.name, override.value, override.textureTransform);
     * });
     * ```
     *
     * @see {@link removeOveride} - To remove a single override
     * @see {@link dispose} - To completely remove the property block and clean up resources
     */
    clearAllOverrides(): void {
        this._overrides = [];
    }

    /**
     * Gets all property overrides as a readonly array
     * @returns Array of all property overrides
     */
    get overrides(): readonly PropertyBlockOverride[] {
        return this._overrides;
    }

    /**
     * Checks if this property block has any overrides
     * @returns True if there are any overrides set
     */
    hasOverrides(): boolean {
        return this._overrides.length > 0;
    }

    /**
     * Set a shader define that will be included in the program cache key.
     * This allows different objects sharing the same material to have different shader programs.
     *
     * Defines affect shader compilation and are useful for enabling/disabling features per-object.
     *
     * @param name The define name (e.g., "USE_LIGHTMAP", "ENABLE_REFLECTIONS")
     * @param value The define value (typically a boolean, number, or string)
     *
     * @example
     * ```typescript
     * // Enable a feature for this specific object
     * block.setDefine("USE_CUSTOM_SHADER", true);
     * block.setDefine("QUALITY_LEVEL", 2);
     * ```
     */
    setDefine(name: string, value: string | number | boolean): void {
        this._defines[name] = value;
    }

    /**
     * Remove a shader define
     * @param name The define name to remove
     */
    clearDefine(name: string): void {
        this._defines[name] = undefined as any;
    }

    /**
     * Get all defines set on this property block
     * @returns A readonly record of all defines
     */
    getDefines(): Readonly<Record<string, string | number | boolean>> {
        return this._defines;
    }

    /**
     * Generates a cache key based on the current overrides and defines.
     * This key is used internally to ensure correct shader program selection
     * when objects share materials but have different property blocks.
     *
     * @returns A string representing the current state of this property block
     * @internal
     */
    getCacheKey(): string {
        const parts: string[] = [];

        // Add defines to cache key
        const defineKeys = Object.keys(this._defines).sort();
        for (const key of defineKeys) {
            const value = this._defines[key];
            if (value !== undefined) {
                parts.push(`d:${key}=${value}`);
            }
        }

        // Add overrides to cache key
        for (const o of this._overrides) {
            if (o.value === null) continue;
            let val = "";
            if (o.value instanceof Texture) {
                val = (o.value as any).uuid || "texture";
                if (o.textureTransform) {
                    const t = o.textureTransform;
                    if (t.offset) val += `;to:${t.offset.x},${t.offset.y}`;
                    if (t.repeat) val += `;tr:${t.repeat.x},${t.repeat.y}`;
                }
            } else if (Array.isArray(o.value)) {
                val = o.value.join(",");
            } else if (o.value && typeof o.value === "object" && "r" in o.value) {
                const c = o.value as any;
                val = `${c.r},${c.g},${c.b},${c.a !== undefined ? c.a : ""}`;
            } else if (o.value && typeof o.value === "object" && "x" in o.value) {
                // Vector2, Vector3, Vector4
                const v = o.value as any;
                val = `${v.x},${v.y}${v.z !== undefined ? `,${v.z}` : ""}${v.w !== undefined ? `,${v.w}` : ""}`;
            } else {
                val = String(o.value);
            }
            parts.push(`${o.name}=${val}`);
        }
        return parts.join(";");
    }
}

/**
 * Symbol used to store original material values on the material object
 * @internal
 */
const $originalValues = Symbol("originalValues");

/**
 * Stores an original material property value before override
 * @internal
 */
interface OriginalValue {
    name: string;
    value: unknown;
}

/**
 * Stores saved texture transform state for restoration
 * @internal
 */
interface SavedTextureTransform {
    name: string;
    offsetX: number;
    offsetY: number;
    repeatX: number;
    repeatY: number;
}

/**
 * Symbol used to store saved texture transforms on the material object
 * @internal
 */
const $savedTextureTransforms = Symbol("savedTextureTransforms");

/**
 * Type for Three.js object render callbacks
 * @internal
 */
type ObjectRenderCallback = (this: Object3D, renderer: WebGLRenderer, scene: Scene, camera: Camera, geometry: BufferGeometry, material: Material, group: Group) => void;

/**
 * Collect all materials from an object and its children
 * @internal
 */
function collectMaterials(object: Object3D, materials: Set<Material>): void {
    const obj = object as Object3D & { material?: Material | Material[] };
    if (obj.material) {
        if (Array.isArray(obj.material)) {
            obj.material.forEach(mat => materials.add(mat));
        } else {
            materials.add(obj.material);
        }
    }

    // For Groups, collect materials from children too
    if (object.type === "Group") {
        object.children.forEach(child => collectMaterials(child, materials));
    }
}

/**
 * Find property block by checking this object and parent (if parent is a Group).
 * Returns both the block and the owner object.
 *
 * @param obj The object to search from
 * @returns The property block and its owner object, or undefined if not found
 * @internal
 */
function findPropertyBlockAndOwner(obj: Object3D): { block: MaterialPropertyBlock; owner: Object3D } | undefined {
    // First check if this object itself has a property block
    let block = registry.getBlock(obj);
    if (block) return { block, owner: obj };

    // If not, check if parent is a Group and has a property block
    if (obj.parent && obj.parent.type === "Group") {
        block = registry.getBlock(obj.parent);
        if (block) return { block, owner: obj.parent };
    }

    return undefined;
}

/**
 * Symbol to track which materials are currently being rendered for an object
 * @internal
 */
const currentlyRenderingFlag = Symbol("beforeRenderingFlag");

/**
 * Tracks original transparent values that were changed during render list building
 * @internal
 */
const beforeRenderListTransparentChanged = new WeakMap<Object3D, boolean>();

/**
 * Tracks original transmission values that were changed during render list building
 * @internal
 */
const beforeRenderListTransparentChangedTransmission = new WeakMap<Object3D, number>();

/**
 * Callback invoked before an object is added to the render list.
 * Used to temporarily override transparency/transmission for correct render list assignment.
 * @internal
 */
const onBeforeRenderListPush = function (this: Object3D, _object: Object3D, _geometry: BufferGeometry, material: Material, _group: Group) {
    const block = registry.getBlock(_object);
    if (!block) {
        return;
    }
    if (block.hasOverrides()) {
        const transmission = block.getOverride<number>("transmission")?.value;
        const transparent = block.getOverride("transparent")?.value;

        if (transmission !== undefined && typeof transmission === "number" && "transmission" in material && transmission !== material.transmission) {
            beforeRenderListTransparentChangedTransmission.set(this, material.transmission as number);
            material.transmission = transmission;
        }
        if (transparent !== undefined && typeof transparent === "boolean" && transparent !== material.transparent) {
            beforeRenderListTransparentChanged.set(this, material.transparent);
            material.transparent = transparent;
        }
    }
}
/**
 * Callback invoked after an object is added to the render list.
 * Restores the original transparency/transmission values that were overridden in onBeforeRenderListPush.
 * @internal
 */
const onAfterRenderListPush = function (this: Object3D, _object: Object3D, _geometry: BufferGeometry, material: Material, _group: Group) {
    const prevTransparent = beforeRenderListTransparentChanged.get(_object);
    if (prevTransparent !== undefined) {
        beforeRenderListTransparentChanged.delete(_object);
        material.transparent = prevTransparent;
    }

    const prevTransmission = beforeRenderListTransparentChangedTransmission.get(_object);
    if (prevTransmission !== undefined) {
        beforeRenderListTransparentChangedTransmission.delete(_object);
        (material as any).transmission = prevTransmission;
    }
}

// #region OnBeforeRender
/**
 * Main callback invoked before rendering an object.
 * Applies property block overrides and defines to the material.
 * @internal
 */
const onBeforeRender_MaterialBlock: ObjectRenderCallback = function (this: Object3D, _renderer: WebGLRenderer, _scene: Scene, _camera: Camera, _geometry: BufferGeometry, material: Material, _group: Group) {


    // Only run if the material belongs to this object and is a "regular" material (not depth or other override material)
    const materials = (this as any).material as Array<Material> | Material | undefined;
    if (!materials) return;
    if (Array.isArray(materials)) {
        if (!materials.includes(material)) return;
    }
    else if (materials !== material) {
        return;
    }

    // Keep track of which materials rendering started for so we can check in onAfterRender if it was processed 
    // (in case of override materials like depth material where onBeforeRender runs but we don't want to apply overrides)
    if (this[currentlyRenderingFlag] === undefined) this[currentlyRenderingFlag] = new WeakSet<Material>();
    this[currentlyRenderingFlag].add(material);


    // Before rendering, check if this object (or its parent Group) has a property block with overrides for this material.
    const result = findPropertyBlockAndOwner(this);
    if (!result) {
        return;
    }

    const { block: propertyBlock, owner } = result;

    // Only apply if this material was registered with this property block
    // if (!registry.hasMaterial(owner, material)) {
    //     return;
    // }

    const overrides = propertyBlock.overrides;
    const mat = material as any;

    // Apply defines to material - this affects shader compilation
    const defines = propertyBlock.getDefines();
    const defineKeys = Object.keys(defines);
    if (defineKeys.length > 0) {
        if (!mat.defines) mat.defines = {};
        for (const key of defineKeys) {
            const value = defines[key];
            if (value !== undefined) {
                mat.defines[key] = value;
            }
        }
    }

    // Still set up cache key even if no overrides (defines affect it)
    if (overrides.length === 0 && defineKeys.length === 0) {
        return;
    }

    // Defines always affect shader compilation → need program change
    let needsProgramChange = defineKeys.length > 0;

    if (!mat[$originalValues]) {
        mat[$originalValues] = [];
    }
    const originalValues = mat[$originalValues] as OriginalValue[];

    for (const override of overrides) {
        if (override.value === null) continue;

        const currentValue = mat[override.name];

        const existingOriginal = originalValues.find((o: OriginalValue) => o.name === override.name);
        if (existingOriginal) {
            // Update to current value each frame so animations/external changes are preserved
            existingOriginal.value = currentValue;
        } else {
            originalValues.push({ name: override.name, value: currentValue });
        }

        // Check if this override changes shader features (truthiness change).
        // E.g. null → Texture enables USE_LIGHTMAP, Texture → null disables it.
        // Pure uniform changes (red → blue, textureA → textureB) don't need program switch.
        if (!needsProgramChange && !!currentValue !== !!override.value) {
            needsProgramChange = true;
        }

        // Set all material properties including lightMap -
        // three.js reads material.lightMap to determine shader parameters and upload uniforms
        mat[override.name] = override.value;

        // Apply per-object texture transform (offset/repeat) if specified
        if (override.textureTransform && override.value instanceof Texture) {
            const tex = override.value;
            if (!mat[$savedTextureTransforms]) mat[$savedTextureTransforms] = [];
            (mat[$savedTextureTransforms] as SavedTextureTransform[]).push({
                name: override.name,
                offsetX: tex.offset.x, offsetY: tex.offset.y,
                repeatX: tex.repeat.x, repeatY: tex.repeat.y
            });
            const t = override.textureTransform;
            if (t.offset) tex.offset.copy(t.offset);
            if (t.repeat) tex.repeat.copy(t.repeat);
        }
    }

    // Only set needsUpdate when overrides change shader features (truthiness changes
    // like null↔texture, or defines added). This triggers getProgram() for program switches.
    // Pure uniform overrides (color, roughness) skip this — no version increment needed.
    if (needsProgramChange) {
        mat.needsUpdate = true;
    }
    // _forceRefresh triggers uniform re-upload for consecutive objects sharing
    // the same program and material (without it three.js skips the upload).
    mat._forceRefresh = true;
};

// #region OnAfterRender
/**
 * Main callback invoked after rendering an object.
 * Restores the original material property values and defines.
 * @internal
 */
const onAfterRender_MaterialBlock: ObjectRenderCallback = function (this: Object3D, _renderer: WebGLRenderer, _scene: Scene, _camera: Camera, _geometry: BufferGeometry, material: Material, _group: Group) {

    // We don't want to run this logic if onBeforeRender didn't run for this material (e.g. due to DepthMaterial or other override material), so we check the flag set in onBeforeRender
    if (this[currentlyRenderingFlag] === undefined) return;
    if (!this[currentlyRenderingFlag].has(material)) return;
    this[currentlyRenderingFlag].delete(material);


    const result = findPropertyBlockAndOwner(this);
    if (!result) {
        return;
    }

    const { block: propertyBlock, owner } = result;

    // Only restore if this material was registered with this property block
    // if (!registry.hasMaterial(owner, material)) {
    //     return;
    // }

    const overrides = propertyBlock.overrides;

    const mat = material as any;
    const originalValues = mat[$originalValues] as OriginalValue[] | undefined;

    // Clean up defines — this affects shader compilation
    const defines = propertyBlock.getDefines();
    const defineKeys = Object.keys(defines);
    let needsProgramChange = false;
    if (defineKeys.length > 0 && mat.defines) {
        for (const key of defineKeys) {
            delete mat.defines[key];
        }
        needsProgramChange = true;
    }

    if (overrides.length === 0) {
        if (needsProgramChange) {
            mat.needsUpdate = true;
            mat._forceRefresh = true;
        }
        return;
    }

    if (!originalValues) return;

    // Restore texture transforms before restoring material properties
    const savedTransforms = mat[$savedTextureTransforms] as SavedTextureTransform[] | undefined;
    if (savedTransforms && savedTransforms.length > 0) {
        for (const saved of savedTransforms) {
            const override = overrides.find(o => o.name === saved.name);
            if (override?.value instanceof Texture) {
                override.value.offset.set(saved.offsetX, saved.offsetY);
                override.value.repeat.set(saved.repeatX, saved.repeatY);
            }
        }
        savedTransforms.length = 0;
    }

    for (const override of overrides) {
        const original = originalValues.find(o => o.name === override.name);
        if (original) {
            // Check if restoring changes shader features (truthiness change)
            if (!needsProgramChange && !!override.value !== !!original.value) {
                needsProgramChange = true;
            }
            mat[override.name] = original.value;
        }
    }

    // Only set needsUpdate when restoring affects shader features
    if (needsProgramChange) {
        mat.needsUpdate = true;
    }
    // Always force uniform refresh so the next object gets correct values
    mat._forceRefresh = true;
};


// #region Attach Callbacks
/**
 * Attaches the property block render callbacks to an object and its child meshes.
 * @param object The object to attach callbacks to
 * @param _propertyBlock The property block being attached (unused but kept for clarity)
 * @internal
 */
function attachPropertyBlockToObject(object: Object3D, _propertyBlock: MaterialPropertyBlock): void {
    // Collect and register all materials that belong to this property block
    // const materials = new Set<Material>();
    // collectMaterials(object, materials);
    // materials.forEach(mat => registry.addMaterial(object, mat));

    // Attach callbacks to renderable objects (Mesh, SkinnedMesh)
    // Groups don't render themselves but we still need to handle child meshes
    if (object.type === "Group") {
        object.children.forEach(child => {
            if (child.type === "Mesh" || child.type === "SkinnedMesh") {
                attachCallbacksToMesh(child, object, _propertyBlock);
            }
        });
    } else if (object.type === "Mesh" || object.type === "SkinnedMesh") {
        attachCallbacksToMesh(object, object, _propertyBlock);
    }
}
/**
 * Attaches render callbacks to a specific mesh object.
 * Chains with existing callbacks if they exist.
 * @param mesh The mesh to attach callbacks to
 * @param propertyBlockOwner The object that owns the property block (may be the mesh itself or its parent Group)
 * @internal
 */
function attachCallbacksToMesh(mesh: Object3D, propertyBlockOwner: Object3D, _propertyBlock: MaterialPropertyBlock): void {
    // Check if this specific mesh already has our callbacks attached for this property block owner
    if (registry.isHooked(mesh, propertyBlockOwner)) {
        // Already hooked for this property block owner
        return;
    }

    registry.addHook(mesh, propertyBlockOwner);

    /**
     * Expose the property block for e.g. Needle Inspector
     */
    mesh["needle:materialPropertyBlock"] = _propertyBlock;

    if (!mesh.onBeforeRender) {
        mesh.onBeforeRender = onBeforeRender_MaterialBlock;
    } else {
        const original = mesh.onBeforeRender;
        mesh.onBeforeRender = function (renderer, scene, camera, geometry, material, group) {
            original.call(this, renderer, scene, camera, geometry, material, group);
            onBeforeRender_MaterialBlock.call(this, renderer, scene, camera, geometry, material, group);
        };
    }

    if (!mesh.onAfterRender) {
        mesh.onAfterRender = onAfterRender_MaterialBlock;
    } else {
        const original = mesh.onAfterRender;
        mesh.onAfterRender = function (renderer, scene, camera, geometry, material, group) {
            onAfterRender_MaterialBlock.call(this, renderer, scene, camera, geometry, material, group);
            original.call(this, renderer, scene, camera, geometry, material, group);
        };
    }

    /** @ts-ignore patched in three.js */
    mesh.onBeforeRenderListPush = onBeforeRenderListPush;
    /** @ts-ignore patched in three.js */
    mesh.onAfterRenderListPush = onAfterRenderListPush;

}
//#endregion

/**
 * Checks if an object has a MaterialPropertyBlock attached to it.
 *
 * @param object The object to check
 * @returns True if the object has a property block registered
 *
 * @example
 * ```typescript
 * if (objectHasPropertyBlock(myMesh)) {
 *   console.log("This mesh has property overrides");
 * }
 * ```
 */
export function objectHasPropertyBlock(object: Object3D): boolean {
    return registry.getBlock(object) !== undefined;
}