popiart-node-canvas/src/utils/browserFileSystem.ts

"use client";

import type { WorkflowFile } from "@/store/workflowStore";

const PREFIX = "browserfs://";
const DB_NAME = "popiart-browser-file-system";
const DB_VERSION = 1;
const STORE_NAME = "directory-handles";
const MAX_WORKFLOW_SCAN_DEPTH = 3;
const SKIP_DIRS = new Set([".git", "node_modules", "__pycache__", ".next"]);

export interface BrowserWorkflowListEntry {
  name: string;
  directoryPath: string;
  relativePath: string;
  lastModified: number;
}

type BrowserDirectoryRecord = {
  id: string;
  handle: FileSystemDirectoryHandle;
};

type PermissionMode = "read" | "readwrite";

type FileSystemPermissionDescriptor = {
  mode?: PermissionMode;
};

type PermissionedDirectoryHandle = FileSystemDirectoryHandle & {
  queryPermission?: (descriptor?: FileSystemPermissionDescriptor) => Promise<PermissionState>;
  requestPermission?: (descriptor?: FileSystemPermissionDescriptor) => Promise<PermissionState>;
};

declare global {
  interface Window {
    showDirectoryPicker?: () => Promise<FileSystemDirectoryHandle>;
  }
}

const memoryHandles = new Map<string, FileSystemDirectoryHandle>();

export function isBrowserFileSystemPath(path: string | null | undefined): boolean {
  return typeof path === "string" && path.startsWith(PREFIX);
}

export function isBrowserDirectoryPickerSupported(): boolean {
  return typeof window !== "undefined" && typeof window.showDirectoryPicker === "function";
}

export function joinBrowserFileSystemPath(basePath: string, folderName: string): string {
  return `${basePath.replace(/\/+$/, "")}/${encodeURIComponent(folderName)}`;
}

const IMAGE_EXTENSIONS = ["png", "jpg", "jpeg", "webp", "gif"] as const;

function extensionFromDataUrl(dataUrl: string): string {
  const mime = dataUrl.match(/^data:([^;]+);base64,/)?.[1];
  if (mime === "image/jpeg") return "jpg";
  if (mime === "image/webp") return "webp";
  if (mime === "image/gif") return "gif";
  return "png";
}

async function dataUrlToBlob(dataUrl: string): Promise<Blob> {
  const response = await fetch(dataUrl);
  return response.blob();
}

async function fileToDataUrl(file: File): Promise<string> {
  return new Promise((resolve, reject) => {
    const reader = new FileReader();
    reader.onload = () => resolve(String(reader.result));
    reader.onerror = () => reject(reader.error);
    reader.readAsDataURL(file);
  });
}

export async function pickBrowserWorkflowDirectory(): Promise<string | null> {
  if (!isBrowserDirectoryPickerSupported() || !window.showDirectoryPicker) return null;

  const handle = await window.showDirectoryPicker();
  await ensurePermission(handle, "readwrite");
  const id = `${Date.now()}-${Math.random().toString(36).slice(2, 10)}`;
  memoryHandles.set(id, handle);
  await putDirectoryRecord({ id, handle });
  return `${PREFIX}${id}`;
}

export async function writeBrowserWorkflowFile(
  directoryPath: string,
  filename: string,
  workflow: WorkflowFile
): Promise<{ filePath: string }> {
  const directoryHandle = await resolveDirectoryHandle(directoryPath, true);
  await ensurePermission(directoryHandle, "readwrite");

  await directoryHandle.getDirectoryHandle("inputs", { create: true });
  await directoryHandle.getDirectoryHandle("generations", { create: true });

  const safeName = filename.replace(/[^a-zA-Z0-9-_]/g, "_");
  const fileName = `${safeName}.json`;
  const fileHandle = await directoryHandle.getFileHandle(fileName, { create: true });
  const writable = await fileHandle.createWritable();
  await writable.write(JSON.stringify(workflow, null, 2));
  await writable.close();

  return { filePath: `${directoryPath}/${fileName}` };
}

export async function writeBrowserGenerationFile(
  directoryPath: string,
  imageId: string,
  imageData: string
): Promise<{ filePath: string }> {
  const directoryHandle = await resolveDirectoryHandle(directoryPath, true);
  await ensurePermission(directoryHandle, "readwrite");

  const extension = extensionFromDataUrl(imageData);
  const fileName = `${imageId}.${extension}`;
  const fileHandle = await directoryHandle.getFileHandle(fileName, { create: true });
  const writable = await fileHandle.createWritable();
  await writable.write(await dataUrlToBlob(imageData));
  await writable.close();

  return { filePath: `${directoryPath}/${fileName}` };
}

export async function loadBrowserGenerationFile(
  directoryPath: string,
  imageId: string
): Promise<string | null> {
  const directoryHandle = await resolveDirectoryHandle(directoryPath, false);
  await ensurePermission(directoryHandle, "readwrite");

  for (const extension of IMAGE_EXTENSIONS) {
    try {
      const fileHandle = await directoryHandle.getFileHandle(`${imageId}.${extension}`);
      return fileToDataUrl(await fileHandle.getFile());
    } catch {
      continue;
    }
  }

  return null;
}

export async function listBrowserWorkflows(parentPath: string): Promise<BrowserWorkflowListEntry[]> {
  const rootHandle = await resolveDirectoryHandle(parentPath, false);
  await ensurePermission(rootHandle, "readwrite");

  const dirs = await collectBrowserDirectories(parentPath, rootHandle, "", 0);
  const results = await Promise.all(
    dirs.map((dir) => probeBrowserWorkflow(dir.path, dir.handle, dir.relativePath))
  );

  return results
    .filter((entry): entry is BrowserWorkflowListEntry => entry !== null)
    .sort((a, b) => b.lastModified - a.lastModified);
}

export async function loadBrowserWorkflowFromDirectory(directoryPath: string): Promise<WorkflowFile> {
  const directoryHandle = await resolveDirectoryHandle(directoryPath, false);
  await ensurePermission(directoryHandle, "readwrite");

  const jsonFiles: { name: string; file: File }[] = [];
  for await (const [name, handle] of directoryEntries(directoryHandle)) {
    if (handle.kind !== "file" || !name.endsWith(".json")) continue;
    const file = await (handle as FileSystemFileHandle).getFile();
    jsonFiles.push({ name, file });
  }

  jsonFiles.sort((a, b) => b.file.lastModified - a.file.lastModified);

  for (const { file } of jsonFiles) {
    try {
      const parsed = JSON.parse(await file.text()) as WorkflowFile;
      if (
        typeof parsed.version === "number" &&
        Array.isArray(parsed.nodes) &&
        Array.isArray(parsed.edges)
      ) {
        return parsed;
      }
    } catch {
      continue;
    }
  }

  throw new Error("No workflow file found in directory");
}

async function collectBrowserDirectories(
  rootPath: string,
  currentHandle: FileSystemDirectoryHandle,
  relativePath: string,
  depth: number
): Promise<Array<{ path: string; handle: FileSystemDirectoryHandle; relativePath: string }>> {
  if (depth > MAX_WORKFLOW_SCAN_DEPTH) return [];

  const dirs: Array<{ path: string; handle: FileSystemDirectoryHandle; relativePath: string }> = [];
  for await (const [name, handle] of directoryEntries(currentHandle)) {
    if (handle.kind !== "directory") continue;
    if (name.startsWith(".") || SKIP_DIRS.has(name)) continue;

    const childHandle = handle as FileSystemDirectoryHandle;
    const childRelativePath = relativePath ? `${relativePath}/${name}` : name;
    const childPath = joinBrowserFileSystemPath(rootPath, childRelativePath);
    dirs.push({ path: childPath, handle: childHandle, relativePath: childRelativePath });

    if (depth < MAX_WORKFLOW_SCAN_DEPTH) {
      try {
        dirs.push(
          ...(await collectBrowserDirectories(rootPath, childHandle, childRelativePath, depth + 1))
        );
      } catch {
        continue;
      }
    }
  }

  return dirs;
}

async function probeBrowserWorkflow(
  directoryPath: string,
  directoryHandle: FileSystemDirectoryHandle,
  relativePath: string
): Promise<BrowserWorkflowListEntry | null> {
  const jsonFiles: { file: File }[] = [];
  for await (const [name, handle] of directoryEntries(directoryHandle)) {
    if (handle.kind !== "file" || !name.endsWith(".json")) continue;
    try {
      jsonFiles.push({ file: await (handle as FileSystemFileHandle).getFile() });
    } catch {
      continue;
    }
  }

  jsonFiles.sort((a, b) => b.file.lastModified - a.file.lastModified);

  for (const { file } of jsonFiles) {
    try {
      const parsed = JSON.parse(await file.text()) as Partial<WorkflowFile>;
      if (
        typeof parsed.version === "number" &&
        Array.isArray(parsed.nodes) &&
        Array.isArray(parsed.edges)
      ) {
        return {
          name: typeof parsed.name === "string" && parsed.name ? parsed.name : dirBasename(relativePath),
          directoryPath,
          relativePath,
          lastModified: file.lastModified,
        };
      }
    } catch {
      continue;
    }
  }

  return null;
}

function dirBasename(dirPath: string): string {
  return dirPath.split("/").filter(Boolean).pop() || dirPath;
}

function directoryEntries(
  handle: FileSystemDirectoryHandle
): AsyncIterable<[string, FileSystemHandle]> {
  return (handle as FileSystemDirectoryHandle & {
    entries: () => AsyncIterable<[string, FileSystemHandle]>;
  }).entries();
}

async function resolveDirectoryHandle(path: string, create: boolean): Promise<FileSystemDirectoryHandle> {
  const parsed = parseBrowserFileSystemPath(path);
  let handle = await getDirectoryHandleById(parsed.id);
  await ensurePermission(handle, "readwrite");

  for (const segment of parsed.segments) {
    handle = await handle.getDirectoryHandle(segment, { create });
  }

  return handle;
}

function parseBrowserFileSystemPath(path: string): { id: string; segments: string[] } {
  if (!isBrowserFileSystemPath(path)) throw new Error("Not a browser file-system path");

  const parts = path
    .slice(PREFIX.length)
    .split("/")
    .filter(Boolean)
    .map((part) => decodeURIComponent(part));

  const id = parts[0];
  if (!id) throw new Error("Browser directory handle is missing");

  return { id, segments: parts.slice(1) };
}

async function getDirectoryHandleById(id: string): Promise<FileSystemDirectoryHandle> {
  const cached = memoryHandles.get(id);
  if (cached) return cached;

  const record = await getDirectoryRecord(id);
  if (!record) {
    throw new Error("Browser folder permission was lost. Choose the folder again.");
  }

  memoryHandles.set(id, record.handle);
  return record.handle;
}

async function ensurePermission(handle: FileSystemDirectoryHandle, mode: PermissionMode): Promise<void> {
  const permissioned = handle as PermissionedDirectoryHandle;
  if (!permissioned.queryPermission || !permissioned.requestPermission) return;

  const descriptor = { mode };
  const current = await permissioned.queryPermission(descriptor);
  if (current === "granted") return;

  const requested = await permissioned.requestPermission(descriptor);
  if (requested !== "granted") {
    throw new Error("Folder permission was not granted.");
  }
}

function openDb(): Promise<IDBDatabase> {
  return new Promise((resolve, reject) => {
    const request = indexedDB.open(DB_NAME, DB_VERSION);
    request.onupgradeneeded = () => {
      request.result.createObjectStore(STORE_NAME, { keyPath: "id" });
    };
    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
  });
}

async function putDirectoryRecord(record: BrowserDirectoryRecord): Promise<void> {
  const db = await openDb();
  await new Promise<void>((resolve, reject) => {
    const tx = db.transaction(STORE_NAME, "readwrite");
    tx.objectStore(STORE_NAME).put(record);
    tx.oncomplete = () => resolve();
    tx.onerror = () => reject(tx.error);
  });
  db.close();
}

async function getDirectoryRecord(id: string): Promise<BrowserDirectoryRecord | null> {
  const db = await openDb();
  const record = await new Promise<BrowserDirectoryRecord | null>((resolve, reject) => {
    const tx = db.transaction(STORE_NAME, "readonly");
    const request = tx.objectStore(STORE_NAME).get(id);
    request.onsuccess = () => resolve((request.result as BrowserDirectoryRecord | undefined) ?? null);
    request.onerror = () => reject(request.error);
  });
  db.close();
  return record;
}