import type { ErrorCode } from "./error";

export const NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION = 1 as const;
export const NANOVDB_PAGE_FEEDBACK_HEADER_WORDS = 4;
export const NANOVDB_PAGE_FEEDBACK_DEFAULT_CAPACITY = 1024;
export const NANOVDB_PAGE_FEEDBACK_MAX_CAPACITY = 8192;
export const NANOVDB_PAGE_FEEDBACK_EMPTY_PAGE_ID = 0xffffffff;
export const NANOVDB_PAGE_FEEDBACK_OVERFLOW_CODE = "NANOVDB_PAGE_FEEDBACK_OVERFLOW" as const satisfies ErrorCode;

export const NANOVDB_PAGE_FEEDBACK_WORD = {
  schemaVersion: 0,
  capacity: 1,
  count: 2,
  overflow: 3,
  pageIds: NANOVDB_PAGE_FEEDBACK_HEADER_WORDS,
} as const;

export const NANOVDB_PAGE_FEEDBACK_WGSL = /* wgsl */`
struct NanoVDBPageFeedback {
  schema_version: u32,
  capacity: u32,
  count: atomic<u32>,
  overflow: atomic<u32>,
  page_ids: array<atomic<u32>>,
}
`;

export const NANOVDB_PAGE_FEEDBACK_RECORD_WGSL = /* wgsl */`
fn nanovdb_record_page_fault(page_id: u32) {
  if (page_id == 0xffffffffu || nanovdb_page_feedback.schema_version != 1u) { return; }
  let physical_capacity = arrayLength(&nanovdb_page_feedback.page_ids);
  let capacity = min(nanovdb_page_feedback.capacity, physical_capacity);
  if (capacity == 0u) { return; }
  for (var slot = 0u; slot < capacity; slot += 1u) {
    loop {
      let current = atomicLoad(&nanovdb_page_feedback.page_ids[slot]);
      if (current == page_id) { return; }
      if (current != 0xffffffffu) { break; }
      let claim = atomicCompareExchangeWeak(&nanovdb_page_feedback.page_ids[slot], 0xffffffffu, page_id);
      if (claim.exchanged) {
        atomicAdd(&nanovdb_page_feedback.count, 1u);
        return;
      }
      if (claim.old_value == page_id) { return; }
      if (claim.old_value != 0xffffffffu) { break; }
    }
  }
  atomicStore(&nanovdb_page_feedback.overflow, 1u);
  atomicAdd(&nanovdb_page_feedback.count, 1u);
}
`;

export type NanoVDBPageFeedbackStatus = "READY" | "OVERFLOW";

export interface NanoVDBPageFeedbackResult {
  schemaVersion: typeof NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION;
  capacity: number;
  attemptedCount: number;
  storedCount: number;
  pageIds: number[];
  status: NanoVDBPageFeedbackStatus;
  errorCode: typeof NANOVDB_PAGE_FEEDBACK_OVERFLOW_CODE | null;
}

export interface NanoVDBPageFeedbackBatch {
  schemaVersion: typeof NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION;
  renderRevision: number;
  attemptedCount: number;
  gpuStoredCount: number;
  uniqueCount: number;
  pageIds: number[];
  status: NanoVDBPageFeedbackStatus;
  errorCode: typeof NANOVDB_PAGE_FEEDBACK_OVERFLOW_CODE | null;
}

export interface NanoVDBPageFeedbackDispatchResult {
  schemaVersion: typeof NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION;
  renderRevision: number;
  currentRenderRevision: number;
  status: "ACCEPTED" | "STALE";
  requestedPageIds: number[];
  requestedCount: number;
  errorCode: "REVISION_CONFLICT" | null;
}

export type NanoVDBPageRequester = (pageId: number, renderRevision: number) => Promise<void> | void;

export class NanoVDBPageFeedbackError extends Error {
  constructor(public readonly code: "INVALID_ARGUMENT" | "PROTOCOL_MISMATCH", message: string) {
    super(`${code}: ${message}`);
    this.name = "NanoVDBPageFeedbackError";
  }
}

function feedbackError(code: NanoVDBPageFeedbackError["code"], message: string): never {
  throw new NanoVDBPageFeedbackError(code, message);
}

function validateCapacity(capacity: number, code: NanoVDBPageFeedbackError["code"] = "INVALID_ARGUMENT"): number {
  if (!Number.isSafeInteger(capacity) || capacity < 1 || capacity > NANOVDB_PAGE_FEEDBACK_MAX_CAPACITY) {
    feedbackError(code, "NanoVDB page feedback capacity is outside the bounded range");
  }
  return capacity;
}

export function nanoVDBPageFeedbackByteLength(capacity = NANOVDB_PAGE_FEEDBACK_DEFAULT_CAPACITY): number {
  return (NANOVDB_PAGE_FEEDBACK_HEADER_WORDS + validateCapacity(capacity)) * Uint32Array.BYTES_PER_ELEMENT;
}

export function createNanoVDBPageFeedbackBuffer(capacity = NANOVDB_PAGE_FEEDBACK_DEFAULT_CAPACITY): ArrayBuffer {
  const words = new Uint32Array(nanoVDBPageFeedbackByteLength(capacity) / Uint32Array.BYTES_PER_ELEMENT);
  words[NANOVDB_PAGE_FEEDBACK_WORD.schemaVersion] = NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION;
  words[NANOVDB_PAGE_FEEDBACK_WORD.capacity] = capacity;
  words.fill(NANOVDB_PAGE_FEEDBACK_EMPTY_PAGE_ID, NANOVDB_PAGE_FEEDBACK_WORD.pageIds);
  return words.buffer;
}

export function resetNanoVDBPageFeedbackBuffer(buffer: ArrayBuffer): void {
  const words = feedbackWords(buffer);
  const capacity = validateHeader(words, buffer.byteLength);
  words[NANOVDB_PAGE_FEEDBACK_WORD.count] = 0;
  words[NANOVDB_PAGE_FEEDBACK_WORD.overflow] = 0;
  words.fill(NANOVDB_PAGE_FEEDBACK_EMPTY_PAGE_ID, NANOVDB_PAGE_FEEDBACK_WORD.pageIds, NANOVDB_PAGE_FEEDBACK_WORD.pageIds + capacity);
}

export function parseNanoVDBPageFeedbackBuffer(buffer: ArrayBuffer, pageCount: number): NanoVDBPageFeedbackResult {
  if (!Number.isSafeInteger(pageCount) || pageCount < 1 || pageCount > NANOVDB_PAGE_FEEDBACK_MAX_CAPACITY) {
    feedbackError("INVALID_ARGUMENT", "NanoVDB virtual page count is outside the bounded range");
  }
  const words = feedbackWords(buffer);
  const capacity = validateHeader(words, buffer.byteLength);
  const attemptedCount = words[NANOVDB_PAGE_FEEDBACK_WORD.count];
  const overflow = words[NANOVDB_PAGE_FEEDBACK_WORD.overflow];
  if (overflow !== 0 && overflow !== 1) feedbackError("PROTOCOL_MISMATCH", "NanoVDB page feedback overflow flag is invalid");
  if ((attemptedCount > capacity) !== (overflow === 1)) {
    feedbackError("PROTOCOL_MISMATCH", "NanoVDB page feedback count and overflow flag disagree");
  }
  const storedCount = Math.min(attemptedCount, capacity);
  const pageIds = [...words.slice(NANOVDB_PAGE_FEEDBACK_WORD.pageIds, NANOVDB_PAGE_FEEDBACK_WORD.pageIds + storedCount)];
  if (pageIds.some((pageId) => pageId === NANOVDB_PAGE_FEEDBACK_EMPTY_PAGE_ID || pageId >= pageCount)) {
    feedbackError("PROTOCOL_MISMATCH", "NanoVDB page feedback contains an invalid virtual page ID");
  }
  const overflowed = overflow === 1;
  return {
    schemaVersion: NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION,
    capacity,
    attemptedCount,
    storedCount,
    pageIds,
    status: overflowed ? "OVERFLOW" : "READY",
    errorCode: overflowed ? NANOVDB_PAGE_FEEDBACK_OVERFLOW_CODE : null,
  };
}

function bindNanoVDBPageFeedbackToRender(
  feedback: NanoVDBPageFeedbackResult,
  renderRevision: number,
): NanoVDBPageFeedbackBatch {
  if (!Number.isSafeInteger(renderRevision) || renderRevision < 0) {
    feedbackError("INVALID_ARGUMENT", "NanoVDB feedback render revision is invalid");
  }
  const pageIds = [...new Set(feedback.pageIds)].sort((left, right) => left - right);
  return {
    schemaVersion: NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION,
    renderRevision,
    attemptedCount: feedback.attemptedCount,
    gpuStoredCount: feedback.storedCount,
    uniqueCount: pageIds.length,
    pageIds,
    status: feedback.status,
    errorCode: feedback.errorCode,
  };
}

export function parseNanoVDBPageFeedbackBatch(
  buffer: ArrayBuffer,
  pageCount: number,
  renderRevision: number,
): NanoVDBPageFeedbackBatch {
  return bindNanoVDBPageFeedbackToRender(parseNanoVDBPageFeedbackBuffer(buffer, pageCount), renderRevision);
}

export async function dispatchNanoVDBPageFeedbackBatch(
  batch: NanoVDBPageFeedbackBatch,
  currentRenderRevision: number,
  requestPage: NanoVDBPageRequester,
): Promise<NanoVDBPageFeedbackDispatchResult> {
  if (!Number.isSafeInteger(currentRenderRevision) || currentRenderRevision < 0) {
    feedbackError("INVALID_ARGUMENT", "NanoVDB current render revision is invalid");
  }
  if (batch.renderRevision !== currentRenderRevision) {
    return {
      schemaVersion: NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION,
      renderRevision: batch.renderRevision,
      currentRenderRevision,
      status: "STALE",
      requestedPageIds: [],
      requestedCount: 0,
      errorCode: "REVISION_CONFLICT",
    };
  }
  const requestedPageIds: number[] = [];
  for (const pageId of batch.pageIds) {
    await requestPage(pageId, batch.renderRevision);
    requestedPageIds.push(pageId);
  }
  return {
    schemaVersion: NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION,
    renderRevision: batch.renderRevision,
    currentRenderRevision,
    status: "ACCEPTED",
    requestedPageIds,
    requestedCount: requestedPageIds.length,
    errorCode: null,
  };
}

function feedbackWords(buffer: ArrayBuffer): Uint32Array {
  if (!(buffer instanceof ArrayBuffer) || buffer.byteLength < NANOVDB_PAGE_FEEDBACK_HEADER_WORDS * Uint32Array.BYTES_PER_ELEMENT || buffer.byteLength % Uint32Array.BYTES_PER_ELEMENT !== 0) {
    feedbackError("PROTOCOL_MISMATCH", "NanoVDB page feedback buffer byte length is invalid");
  }
  return new Uint32Array(buffer);
}

function validateHeader(words: Uint32Array, byteLength: number): number {
  if (words[NANOVDB_PAGE_FEEDBACK_WORD.schemaVersion] !== NANOVDB_PAGE_FEEDBACK_SCHEMA_VERSION) {
    feedbackError("PROTOCOL_MISMATCH", "NanoVDB page feedback schema version is unsupported");
  }
  const capacity = validateCapacity(words[NANOVDB_PAGE_FEEDBACK_WORD.capacity], "PROTOCOL_MISMATCH");
  if (byteLength !== nanoVDBPageFeedbackByteLength(capacity)) {
    feedbackError("PROTOCOL_MISMATCH", "NanoVDB page feedback buffer does not match its declared capacity");
  }
  return capacity;
}