From 0a2336814a421d3593ab54b82bedd2de6b3b61c2 Mon Sep 17 00:00:00 2001
From: Felvesthe <23108389+Felvesthe@users.noreply.github.com>
Date: Fri, 16 May 2025 21:57:28 +0200
Subject: [PATCH] Fix lock icon visibility after folder expansion

Signed-off-by: Felvesthe <23108389+Felvesthe@users.noreply.github.com>
---
 src/main.ts            |  3 +--
 src/ui/fileExplorer.ts | 47 ++++++++++++++++++++++++++++++++++++++----
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/src/main.ts b/src/main.ts
index 029f9ea..0ab3848 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -53,14 +53,13 @@ export default class NoteLockerPlugin extends Plugin {
 				return false;
 			},
 		});
-
 	}
 
 	onunload() {
 		const styleEl = document.getElementById('note-locker-styles');
 		if (styleEl) styleEl.remove();
 
-		this.fileExplorerUI.removeFileExplorerIcons();
+		this.fileExplorerUI.cleanup();
 		this.statusBarUI.removeStatusBarItem();
 	}
 
diff --git a/src/ui/fileExplorer.ts b/src/ui/fileExplorer.ts
index 392fc61..c471a4c 100644
--- a/src/ui/fileExplorer.ts
+++ b/src/ui/fileExplorer.ts
@@ -4,6 +4,7 @@ import type NoteLockerPlugin from "../main";
 export class FileExplorerUI {
 	private plugin: NoteLockerPlugin;
 	private folderObserver: MutationObserver | null = null;
+	private updateDebounceTimeout: number | null = null;
 
 	constructor(plugin: NoteLockerPlugin) {
 		this.plugin = plugin;
@@ -38,10 +39,40 @@ export class FileExplorerUI {
 		document.head.appendChild(styleEl);
 
 		// Initial update
-		setTimeout(() => {
+		this.plugin.app.workspace.onLayoutReady(() => {
 			this.updateFileExplorerIcons();
 			this.setupFolderObserver();
-		}, 500);
+
+			// Register handlers for folder expansion/collapse
+			this.registerFolderClickHandlers();
+		});
+	}
+
+	private registerFolderClickHandlers(): void {
+		const fileExplorer = document.querySelector('.workspace-split.mod-left-split .nav-files-container');
+		if (!fileExplorer) return;
+
+		fileExplorer.addEventListener('click', (e) => {
+			const target = e.target as HTMLElement;
+			const folderCollapseIndicator = target.closest('.nav-folder-collapse-indicator') ||
+				target.closest('.nav-folder-title');
+
+			if (folderCollapseIndicator) {
+				// A folder was clicked, so schedule an update after DOM changes
+				this.debouncedUpdateIcons(250);
+			}
+		});
+	}
+
+	private debouncedUpdateIcons(delay: number = 100): void {
+		if (this.updateDebounceTimeout !== null) {
+			window.clearTimeout(this.updateDebounceTimeout);
+		}
+
+		this.updateDebounceTimeout = window.setTimeout(() => {
+			this.updateFileExplorerIcons();
+			this.updateDebounceTimeout = null;
+		}, delay);
 	}
 
 	private setupFolderObserver(): void {
@@ -72,8 +103,7 @@ export class FileExplorerUI {
 			}
 
 			if (shouldUpdate) {
-				// Avoid too many rapid updates
-				setTimeout(() => this.updateFileExplorerIcons(), 50);
+				this.debouncedUpdateIcons(150);
 			}
 		});
 
@@ -115,10 +145,19 @@ export class FileExplorerUI {
 
 	public removeFileExplorerIcons(): void {
 		document.querySelectorAll('.note-locker-icon').forEach(el => el.remove());
+	}
+
+	public cleanup(): void {
+		if (this.updateDebounceTimeout !== null) {
+			window.clearTimeout(this.updateDebounceTimeout);
+			this.updateDebounceTimeout = null;
+		}
 
 		if (this.folderObserver) {
 			this.folderObserver.disconnect();
 			this.folderObserver = null;
 		}
+
+		this.removeFileExplorerIcons();
 	}
 }