import {
  OOM_RECOVERY_REPORT_SCHEMA,
  parseOOMRecoverySuite,
  type OOMFaultObservationIR,
  type OOMFaultPoint,
  type OOMRecoveryReportIR,
} from "../../../protocol/oom-recovery";
import { WebEngineClient } from "../engine-client/WebEngineClient";
import {
  readProjectBlend,
  recoverProjectBlend,
  writeProjectBlend,
} from "../storage/opfs-files";
import { createNanoVDBFloat32GridPaged } from "../render/nanovdb-volume-renderer";
import {
  BoxGeometry,
  BufferGeometry,
  Float32BufferAttribute,
  Mesh,
  MeshBasicMaterial,
  PerspectiveCamera,
  Scene,
  WebGLRenderer,
} from "../vendor/three/three.module.js";
import {
  assertFault,
  beginOOMFaultSession,
  OOMFaultSessionAccessError,
  type OOMAllocationLease,
  type OOMFaultSession,
} from "./oom-fault-session";

export { parseOOMRecoverySuite } from "../../../protocol/oom-recovery";

type OpfsStorage = StorageManager & { getDirectory?: () => Promise<FileSystemDirectoryHandle> };
type DirectoryEntries = AsyncIterableIterator<[string, FileSystemHandle]>;
type TestGPUBufferDescriptor = Parameters<GPUDevice["createBuffer"]>[0];

async function sha256(value: ArrayBuffer | string): Promise<string> {
  const bytes = typeof value === "string" ? new TextEncoder().encode(value) : new Uint8Array(value);
  const digest = await crypto.subtle.digest("SHA-256", bytes);
  return Array.from(new Uint8Array(digest), (byte) => byte.toString(16).padStart(2, "0")).join("");
}

function fault(session: OOMFaultSession, token: string, point: OOMFaultPoint, bytes: number): OOMFaultObservationIR {
  try {
    session.reserve(token, point, bytes, `${point.toLowerCase()}-failure`);
  }
  catch (error) {
    return assertFault(error, point);
  }
  throw new Error(`OOM_FAULT_NOT_TRIGGERED: ${point}`);
}

function proveTokenIsolation(session: OOMFaultSession, token: string, point: OOMFaultPoint): boolean {
  try {
    session.reserve(`${token}-unauthorized`, point, 1, "unauthorized-allocation");
  }
  catch (error) {
    return error instanceof OOMFaultSessionAccessError;
  }
  return false;
}

async function runWasmMainScenario(input: ArrayBuffer): Promise<OOMRecoveryReportIR> {
  const engine = new WebEngineClient({ timeoutMs: 30_000 });
  const restarted = new WebEngineClient({ timeoutMs: 30_000 });
  const faults: OOMFaultObservationIR[] = [];
  let openIsolation: boolean;
  let editIsolation: boolean;
  let saveIsolation: boolean;
  let nativeBefore: number;
  let nativeAfter: number;
  let saveBytes: number;
  let revisionBefore: number;
  let revisionAfter: number;
  let hashBefore: string;
  let hashAfter: string;
  try {
    const initialized = await engine.init();
    if (!initialized.ready || initialized.liveHandles !== 1) throw new Error("WASM OOM fixture did not initialize one clean handle");

    const open = beginOOMFaultSession({ point: "WASM_MAIN_OPEN_INPUT", failAfterCount: 0 });
    openIsolation = proveTokenIsolation(open.session, open.token, "WASM_MAIN_OPEN_INPUT");
    faults.push(fault(open.session, open.token, "WASM_MAIN_OPEN_INPUT", input.byteLength));
    const openStats = open.session.close(open.token);
    if (!openStats.triggered || openStats.liveResources !== 0) throw new Error("WASM open fault leaked a test allocation");
    const disabledLease = open.session.reserve(open.token, "WASM_MAIN_OPEN_INPUT", input.byteLength, "closed-session-probe");
    if (disabledLease.tracked) throw new Error("Closed OOM session leaked into the next request");

    const opened = await engine.openBlend(input.slice(0));
    nativeBefore = opened.status.allocatedBytes;
    const created = await engine.applyCommand({ type: "createPrimitive", primitive: "CUBE", name: "OOMRecoveryCube", location: [2, 0, 0] });
    const createdObject = created.snapshot.nodes.find((node) => node.name === "OOMRecoveryCube");
    if (!createdObject) throw new Error("WASM OOM fixture could not establish an undoable Main revision");
    revisionBefore = created.snapshot.revision;
    hashBefore = await sha256(JSON.stringify(created.snapshot));

    const edit = beginOOMFaultSession({ point: "WASM_MAIN_EDIT_COMMAND", failAfterBytes: 0 });
    editIsolation = proveTokenIsolation(edit.session, edit.token, "WASM_MAIN_EDIT_COMMAND");
    const editBytes = new TextEncoder().encode(JSON.stringify({ type: "setFrame", frame: 21 })).byteLength;
    faults.push(fault(edit.session, edit.token, "WASM_MAIN_EDIT_COMMAND", editBytes));
    const editStats = edit.session.close(edit.token);
    if (!editStats.triggered || editStats.liveResources !== 0) throw new Error("WASM edit fault leaked a test allocation");
    const unchanged = await engine.snapshot();
    revisionAfter = unchanged.snapshot.revision;
    hashAfter = await sha256(JSON.stringify(unchanged.snapshot));
    if (revisionAfter !== revisionBefore || hashAfter !== hashBefore) throw new Error("WASM edit OOM changed Main state");

    const undone = await engine.applyCommand({ type: "undo" });
    if (undone.snapshot.nodes.some((node) => node.id === createdObject.id)) throw new Error("WASM edit OOM damaged the undo stack");
    const redone = await engine.applyCommand({ type: "redo" });
    if (!redone.snapshot.nodes.some((node) => node.id === createdObject.id)) throw new Error("WASM edit OOM damaged the redo stack");

    const firstSave = await engine.saveBlend();
    saveBytes = firstSave.byteLength;
    const save = beginOOMFaultSession({ point: "WASM_MAIN_SAVE_RESULT", failAfterBytes: Math.max(0, saveBytes - 1) });
    saveIsolation = proveTokenIsolation(save.session, save.token, "WASM_MAIN_SAVE_RESULT");
    faults.push(fault(save.session, save.token, "WASM_MAIN_SAVE_RESULT", saveBytes));
    const saveStats = save.session.close(save.token);
    if (!saveStats.triggered || saveStats.liveResources !== 0) throw new Error("WASM save fault leaked a test allocation");
    const afterSaveFault = await engine.snapshot();
    if (!afterSaveFault.snapshot.nodes.some((node) => node.id === createdObject.id)) throw new Error("WASM save OOM changed Main state");
    nativeAfter = afterSaveFault.status.allocatedBytes;

    const saved = await engine.saveBlend();
    engine.terminate();
    const reopened = await restarted.openBlend(saved);
    if (!reopened.snapshot.nodes.some((node) => node.id === createdObject.id)) throw new Error("WASM OOM restart lost the last valid Main");
  }
  finally {
    engine.terminate();
    restarted.terminate();
  }

  const peakBytes = Math.max(nativeBefore, nativeAfter, nativeBefore + saveBytes);
  return {
    schemaVersion: OOM_RECOVERY_REPORT_SCHEMA,
    scenario: "WASM_MAIN",
    faults,
    memory: { beforeBytes: nativeBefore, peakBytes, afterBytes: nativeAfter, releasedBytes: saveBytes },
    state: { revisionBefore, revisionAfter, hashBefore, hashAfter, temporaryResourcesBefore: 0, temporaryResourcesPeak: 0, temporaryResourcesAfter: 0 },
    recovery: { recovered: true, sameSession: true, restartedSession: true, tokenIsolated: openIsolation && editIsolation && saveIsolation },
    checks: ["open-handle-clean", "main-revision-stable", "undo-redo-stable", "save-buffer-released", "worker-reopen"],
  };
}

function wrapStageOOMFile(
  file: FileSystemFileHandle,
  session: OOMFaultSession,
  token: string,
): FileSystemFileHandle {
  return new Proxy(file, {
    get(target, property) {
      if (property === "createWritable") {
        return async (...args: Parameters<FileSystemFileHandle["createWritable"]>) => {
          const writable = await target.createWritable(...args);
          return new Proxy(writable, {
            get(stream, streamProperty) {
              if (streamProperty === "write") {
                return async (value: FileSystemWriteChunkType) => {
                  if (!(value instanceof ArrayBuffer)) return stream.write(value);
                  const firstLength = Math.max(1, Math.floor(value.byteLength / 2));
                  let lease: OOMAllocationLease | undefined;
                  try {
                    lease = session.reserve(token, "OPFS_STAGING_WRITE", firstLength, "opfs-partial-stage");
                    await stream.write(value.slice(0, firstLength));
                    session.reserve(token, "OPFS_STAGING_WRITE", value.byteLength - firstLength, "opfs-stage-remainder");
                  }
                  catch (error) {
                    await stream.abort(error).catch(() => undefined);
                    lease?.release();
                    throw error;
                  }
                  throw new Error("OPFS OOM injector failed to stop the staging write");
                };
              }
              const current = Reflect.get(stream, streamProperty, stream);
              return typeof current === "function" ? current.bind(stream) : current;
            },
          });
        };
      }
      const current = Reflect.get(target, property, target);
      return typeof current === "function" ? current.bind(target) : current;
    },
  });
}

function wrapStageOOMDirectory(
  directory: FileSystemDirectoryHandle,
  session: OOMFaultSession,
  token: string,
): FileSystemDirectoryHandle {
  return new Proxy(directory, {
    get(target, property) {
      if (property === "getDirectoryHandle") {
        return async (...args: Parameters<FileSystemDirectoryHandle["getDirectoryHandle"]>) =>
          wrapStageOOMDirectory(await target.getDirectoryHandle(...args), session, token);
      }
      if (property === "getFileHandle") {
        return async (...args: Parameters<FileSystemDirectoryHandle["getFileHandle"]>) => {
          const file = await target.getFileHandle(...args);
          return args[0].endsWith(".stage") ? wrapStageOOMFile(file, session, token) : file;
        };
      }
      const current = Reflect.get(target, property, target);
      return typeof current === "function" ? current.bind(target) : current;
    },
  });
}

function stageOOMStorage(session: OOMFaultSession, token: string): StorageManager {
  const storage = navigator.storage as OpfsStorage;
  if (!storage.getDirectory) throw new Error("OPFS is unavailable for OOM recovery");
  return new Proxy(storage, {
    get(target, property) {
      if (property === "getDirectory") return async () => wrapStageOOMDirectory(await target.getDirectory!(), session, token);
      const current = Reflect.get(target, property, target);
      return typeof current === "function" ? current.bind(target) : current;
    },
  });
}

async function projectTemporaryResources(projectId: string): Promise<number> {
  const storage = navigator.storage as OpfsStorage;
  if (!storage.getDirectory) throw new Error("OPFS is unavailable for OOM recovery");
  let directory = await storage.getDirectory();
  for (const name of ["projects", projectId]) directory = await directory.getDirectoryHandle(name);
  let count = 0;
  const entries = (directory as unknown as { entries: () => DirectoryEntries }).entries();
  for await (const [name] of entries) if (name.endsWith(".stage")) count++;
  const tmp = await directory.getDirectoryHandle("tmp");
  const tmpEntries = (tmp as unknown as { entries: () => DirectoryEntries }).entries();
  for await (const [name] of tmpEntries) if (name.endsWith(".journal.json") || name.endsWith(".stage") || name.endsWith(".tmp")) count++;
  return count;
}

async function removeProject(projectId: string): Promise<void> {
  const storage = navigator.storage as OpfsStorage;
  if (!storage.getDirectory) return;
  const root = await storage.getDirectory();
  try { await (await root.getDirectoryHandle("projects")).removeEntry(projectId, { recursive: true }); }
  catch (error) { if (!(error instanceof DOMException) || error.name !== "NotFoundError") throw error; }
}

async function runOpfsScenario(input: ArrayBuffer): Promise<OOMRecoveryReportIR> {
  const projectId = `oom-opfs-${crypto.randomUUID()}`;
  const oldBytes = input.slice(0);
  const newBytes = input.slice(0);
  const view = new Uint8Array(newBytes);
  view[view.byteLength - 1] ^= 0x01;
  const oldHash = await sha256(oldBytes);
  let afterHash: string;
  let temporaryBefore: number;
  let temporaryAfter: number;
  let observation: OOMFaultObservationIR | undefined;
  let isolated: boolean;
  let peakBytes: number;
  try {
    await writeProjectBlend(projectId, 7, oldBytes.slice(0));
    temporaryBefore = await projectTemporaryResources(projectId);
    const started = beginOOMFaultSession({ point: "OPFS_STAGING_WRITE", failAfterBytes: Math.max(1, Math.floor(newBytes.byteLength / 2)) });
    isolated = proveTokenIsolation(started.session, started.token, "OPFS_STAGING_WRITE");
    try {
      await writeProjectBlend(projectId, 8, newBytes, stageOOMStorage(started.session, started.token));
      throw new Error("OPFS staging OOM did not reject the write");
    }
    catch (error) {
      observation = assertFault(error, "OPFS_STAGING_WRITE");
    }
    const stats = started.session.close(started.token);
    peakBytes = stats.peakBytes;
    if (!stats.triggered || stats.currentBytes !== 0 || stats.liveResources !== 0) throw new Error("OPFS staging OOM leaked its partial allocation");

    const stored = await readProjectBlend(projectId);
    afterHash = await sha256(stored);
    const recovered = await recoverProjectBlend(projectId);
    temporaryAfter = await projectTemporaryResources(projectId);
    if (oldHash !== afterHash || recovered.status !== "clean" || recovered.manifest?.revision !== 7) throw new Error("OPFS staging OOM replaced the last committed revision");
    if (temporaryAfter !== 0) throw new Error("OPFS staging OOM left temporary files");
    const next = await writeProjectBlend(projectId, 8, oldBytes.slice(0));
    if (next.manifest.revision !== 8) throw new Error("OPFS did not accept a save after OOM recovery");
  }
  finally {
    await removeProject(projectId);
  }
  if (!observation) throw new Error("OPFS OOM observation is missing");
  return {
    schemaVersion: OOM_RECOVERY_REPORT_SCHEMA,
    scenario: "OPFS_STAGING",
    faults: [observation],
    memory: { beforeBytes: 0, peakBytes, afterBytes: 0, releasedBytes: peakBytes },
    state: { revisionBefore: 7, revisionAfter: 7, hashBefore: oldHash, hashAfter: afterHash, temporaryResourcesBefore: temporaryBefore, temporaryResourcesPeak: Math.max(1, temporaryBefore), temporaryResourcesAfter: temporaryAfter },
    recovery: { recovered: true, sameSession: true, restartedSession: true, tokenIsolated: isolated },
    checks: ["partial-stage-aborted", "journal-not-committed", "old-revision-preserved", "old-hash-preserved", "temporary-files-removed", "next-save-succeeds"],
  };
}

function readPixels(renderer: WebGLRenderer): ArrayBuffer {
  const pixels = new Uint8Array(64 * 64 * 4);
  renderer.getContext().readPixels(0, 0, 64, 64, renderer.getContext().RGBA, renderer.getContext().UNSIGNED_BYTE, pixels);
  return pixels.buffer;
}

async function runGpuScenario(): Promise<OOMRecoveryReportIR> {
  const canvas = document.createElement("canvas");
  const renderer = new WebGLRenderer({ canvas, antialias: false, preserveDrawingBuffer: true });
  renderer.setSize(64, 64, false);
  const scene = new Scene();
  const camera = new PerspectiveCamera(50, 1, 0.1, 100);
  camera.position.z = 4;
  const baselineGeometry = new BoxGeometry(1, 1, 1);
  const baselineMaterial = new MeshBasicMaterial({ color: 0x33aa66 });
  const baseline = new Mesh(baselineGeometry, baselineMaterial);
  scene.add(baseline);
  renderer.render(scene, camera);
  const hashBefore = await sha256(readPixels(renderer));

  const started = beginOOMFaultSession({ point: "GPU_TEXTURE_UPLOAD", failAfterCount: 0 });
  const isolated = proveTokenIsolation(started.session, started.token, "GPU_TEXTURE_UPLOAD");
  let geometry: BufferGeometry | undefined;
  let material: MeshBasicMaterial | undefined;
  let geometryLease: OOMAllocationLease | undefined;
  let observation: OOMFaultObservationIR | undefined;
  try {
    const positions = new Float32Array([-0.5, -0.5, 0, 0.5, -0.5, 0, 0, 0.5, 0]);
    geometryLease = started.session.reserve(started.token, "GPU_GEOMETRY_UPLOAD", positions.byteLength, "gpu-partial-geometry");
    geometry = new BufferGeometry();
    geometry.setAttribute("position", new Float32BufferAttribute(positions, 3));
    material = new MeshBasicMaterial({ color: 0xffffff });
    started.session.reserve(started.token, "GPU_TEXTURE_UPLOAD", 4 * 1024 * 1024, "gpu-failed-texture");
    scene.add(new Mesh(geometry, material));
    throw new Error("GPU texture OOM did not stop the atomic scene object");
  }
  catch (error) {
    observation = assertFault(error, "GPU_TEXTURE_UPLOAD");
    geometry?.dispose();
    material?.dispose();
    geometryLease?.release();
  }
  const stats = started.session.close(started.token);
  if (!stats.triggered || stats.liveResources !== 0) throw new Error("GPU OOM leaked a tracked resource");
  renderer.render(scene, camera);
  const hashAfter = await sha256(readPixels(renderer));
  if (hashAfter !== hashBefore || scene.children.length !== 1) throw new Error("GPU OOM published a partial scene object");

  const recoveryGeometry = new BoxGeometry(0.5, 0.5, 0.5);
  const recoveryMaterial = new MeshBasicMaterial({ color: 0xcc3355 });
  const recoveryMesh = new Mesh(recoveryGeometry, recoveryMaterial);
  scene.add(recoveryMesh);
  renderer.render(scene, camera);
  const recoveryPixels = new Uint8Array(readPixels(renderer));
  if (!recoveryPixels.some((value, index) => index % 4 !== 3 && value > 0)) throw new Error("GPU did not render after OOM cleanup");
  scene.remove(recoveryMesh);
  recoveryGeometry.dispose();
  recoveryMaterial.dispose();
  scene.remove(baseline);
  baselineGeometry.dispose();
  baselineMaterial.dispose();
  renderer.dispose();
  if (!observation) throw new Error("GPU OOM observation is missing");
  return {
    schemaVersion: OOM_RECOVERY_REPORT_SCHEMA,
    scenario: "GPU_RESOURCES",
    faults: [observation],
    memory: { beforeBytes: 0, peakBytes: stats.peakBytes, afterBytes: stats.currentBytes, releasedBytes: stats.releasedBytes },
    state: { revisionBefore: 0, revisionAfter: 0, hashBefore, hashAfter, temporaryResourcesBefore: 0, temporaryResourcesPeak: stats.peakResources, temporaryResourcesAfter: stats.liveResources },
    recovery: { recovered: true, sameSession: true, restartedSession: false, tokenIsolated: isolated },
    checks: ["partial-geometry-disposed", "partial-object-not-published", "previous-frame-stable", "small-scene-renders"],
  };
}

async function runNanoVdbScenario(): Promise<OOMRecoveryReportIR> {
  if (!navigator.gpu) throw new Error("WebGPU is unavailable for NanoVDB OOM recovery");
  const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
  if (!adapter) throw new Error("WebGPU adapter is unavailable for NanoVDB OOM recovery");
  const device = await adapter.requestDevice();
  const started = beginOOMFaultSession({ point: "NANOVDB_PAGE_TABLE", failAfterCount: 0 });
  const isolated = proveTokenIsolation(started.session, started.token, "NANOVDB_PAGE_TABLE");
  const destroyed = new Map<string, number>();
  const trackedDevice = new Proxy(device, {
    get(target, property) {
      if (property === "createBuffer") {
        return (descriptor: TestGPUBufferDescriptor): GPUBuffer => {
          const label = descriptor.label ?? "unlabeled";
          const point: OOMFaultPoint = label.includes("page table") ? "NANOVDB_PAGE_TABLE" : "NANOVDB_RESIDENT_BUFFER";
          const lease = started.session.reserve(started.token, point, Number(descriptor.size), `nanovdb-${label}`);
          const buffer = target.createBuffer(descriptor);
          let released = false;
          return new Proxy(buffer, {
            get(bufferTarget, bufferProperty) {
              if (bufferProperty === "destroy") {
                return () => {
                  if (!released) {
                    released = true;
                    destroyed.set(label, (destroyed.get(label) ?? 0) + 1);
                    lease.release();
                  }
                  bufferTarget.destroy();
                };
              }
              const current = Reflect.get(bufferTarget, bufferProperty, bufferTarget);
              return typeof current === "function" ? current.bind(bufferTarget) : current;
            },
          });
        };
      }
      const current = Reflect.get(target, property, target);
      return typeof current === "function" ? current.bind(target) : current;
    },
  });

  let observation: OOMFaultObservationIR | undefined;
  try {
    createNanoVDBFloat32GridPaged(trackedDevice, 128 * 1024, 64 * 1024, 64 * 1024);
    throw new Error("NanoVDB page-table OOM did not reject the grid");
  }
  catch (error) {
    observation = assertFault(error, "NANOVDB_PAGE_TABLE");
  }
  const stats = started.session.close(started.token);
  if (!stats.triggered || stats.liveResources !== 0 || destroyed.get("NanoVDB paged Float32 grid") !== 1) {
    throw new Error("NanoVDB OOM did not uniquely release the resident buffer");
  }

  const recovered = createNanoVDBFloat32GridPaged(device, 128 * 1024, 64 * 1024, 64 * 1024);
  recovered.uploadPage(0, new ArrayBuffer(64 * 1024));
  recovered.uploadPage(1, new ArrayBuffer(64 * 1024));
  if (!recovered.hasResidentPage(1) || recovered.hasResidentPage(0) || recovered.evictionCount !== 1) throw new Error("NanoVDB resident allocator did not recover on the same device");
  recovered.dispose();
  device.destroy();
  if (!observation) throw new Error("NanoVDB OOM observation is missing");
  return {
    schemaVersion: OOM_RECOVERY_REPORT_SCHEMA,
    scenario: "NANOVDB_RESIDENT",
    faults: [observation],
    memory: { beforeBytes: 0, peakBytes: stats.peakBytes, afterBytes: stats.currentBytes, releasedBytes: stats.releasedBytes },
    state: { revisionBefore: 0, revisionAfter: 0, temporaryResourcesBefore: 0, temporaryResourcesPeak: stats.peakResources, temporaryResourcesAfter: stats.liveResources },
    recovery: { recovered: true, sameSession: true, restartedSession: false, tokenIsolated: isolated },
    checks: ["resident-buffer-destroyed-once", "page-table-not-published", "resident-pages-recreated", "same-device-recovers", "lru-eviction-recovers"],
  };
}

export async function runOOMRecoveryScenarios(input: ArrayBuffer): Promise<OOMRecoveryReportIR[]> {
  if (!(input instanceof ArrayBuffer) || input.byteLength === 0) throw new Error("OOM recovery requires a non-empty .blend fixture");
  return parseOOMRecoverySuite([
    await runWasmMainScenario(input.slice(0)),
    await runOpfsScenario(input.slice(0)),
    await runGpuScenario(),
    await runNanoVdbScenario(),
  ]);
}