import { validateWasmModuleManifest, type WasmModuleEntry, type WasmModuleManifest, type WasmModuleVariant } from "../../../protocol/wasm-modules";

export type WasmModuleLoadStatus = "pending" | "loaded" | "failed";

export interface WasmModuleState {
  id: string;
  lifecycle: WasmModuleEntry["lifecycle"];
  status: WasmModuleLoadStatus;
  byteLength: number;
  error?: string;
}

export interface WasmModuleLoader {
  preloadCompleted(): Promise<WasmModuleState[]>;
  load(id: string): Promise<ArrayBuffer>;
  loadDebug(id: string): Promise<ArrayBuffer>;
  state(): WasmModuleState[];
  manifest(): WasmModuleManifest;
}

async function fetchVerified(resource: { url: string; sha256: string }, fetcher: typeof fetch): Promise<ArrayBuffer> {
  const response = await fetcher(resource.url, { cache: "force-cache" });
  if (!response.ok) throw new Error(`WASM_MODULE_REQUEST_FAILED: ${resource.url}:${response.status}`);
  const bytes = await response.arrayBuffer();
  const digest = await crypto.subtle.digest("SHA-256", bytes);
  const actual = [...new Uint8Array(digest)].map((byte) => byte.toString(16).padStart(2, "0")).join("");
  if (actual !== resource.sha256) throw new Error(`WASM_MODULE_HASH_MISMATCH: ${resource.url}`);
  return bytes;
}

export function createWasmModuleLoader(
  rawManifest: unknown,
  variant: WasmModuleVariant,
  fetcher: typeof fetch = fetch,
): WasmModuleLoader {
  const manifest = validateWasmModuleManifest(rawManifest);
  const bytes = new Map<string, ArrayBuffer>();
  const statuses = new Map<string, WasmModuleState>(manifest.modules.map((entry) => [entry.id, { id: entry.id, lifecycle: entry.lifecycle, status: "pending", byteLength: 0 }]));
  const loading = new Map<string, Promise<ArrayBuffer>>();
  const markLoaded = (entry: WasmModuleEntry, value: ArrayBuffer): ArrayBuffer => {
    statuses.set(entry.id, { id: entry.id, lifecycle: entry.lifecycle, status: "loaded", byteLength: value.byteLength });
    return value;
  };
  const load = (entry: WasmModuleEntry): Promise<ArrayBuffer> => {
    const resource = entry.resources[variant];
    const existing = bytes.get(resource.url);
    if (existing) return Promise.resolve(markLoaded(entry, existing));
    const active = loading.get(resource.url);
    if (active) return active.then((value) => markLoaded(entry, value), (error: unknown) => {
      statuses.set(entry.id, { id: entry.id, lifecycle: entry.lifecycle, status: "failed", byteLength: 0, error: error instanceof Error ? error.message : String(error) });
      throw error;
    });
    const promise = fetchVerified(resource, fetcher).then((value) => {
      bytes.set(resource.url, value);
      return markLoaded(entry, value);
    }).catch((error: unknown) => {
      statuses.set(entry.id, { id: entry.id, lifecycle: entry.lifecycle, status: "failed", byteLength: 0, error: error instanceof Error ? error.message : String(error) });
      throw error;
    }).finally(() => loading.delete(resource.url));
    loading.set(resource.url, promise);
    return promise;
  };
  return {
    async preloadCompleted() {
      await Promise.all(manifest.modules.filter((entry) => entry.lifecycle === "complete" && entry.preload).map(load));
      return [...statuses.values()].map((state) => ({ ...state }));
    },
    load: async (id) => {
      const entry = manifest.modules.find((candidate) => candidate.id === id);
      if (!entry) throw new Error(`WASM_MODULE_NOT_FOUND: ${id}`);
      return load(entry);
    },
    async loadDebug(id) {
      const entry = manifest.modules.find((candidate) => candidate.id === id);
      if (!entry || entry.lifecycle !== "debug") throw new Error(`WASM_MODULE_DEBUG_NOT_FOUND: ${id}`);
      return load(entry);
    },
    state: () => [...statuses.values()].map((state) => ({ ...state })),
    manifest: () => manifest,
  };
}