import { open, readdir, stat, writeFile, mkdir } from "node:fs/promises";
import path from "node:path";
import process from "node:process";
import { fileURLToPath } from "node:url";
import sharp from "sharp";

const scriptDirectory = path.dirname(fileURLToPath(import.meta.url));
const projectRoot = path.resolve(scriptDirectory, "../..");
const KTX2_IDENTIFIER = Buffer.from([0xab, 0x4b, 0x54, 0x58, 0x20, 0x32, 0x30, 0xbb, 0x0d, 0x0a, 0x1a, 0x0a]);

function argumentValue(args, name, fallback) {
  const index = args.indexOf(name);
  return index >= 0 && args[index + 1] ? args[index + 1] : fallback;
}

function positionalArguments(args) {
  const result = [];
  for (let index = 0; index < args.length; index += 1) {
    if (args[index] === "--report") {
      index += 1;
      continue;
    }
    if (!args[index].startsWith("--")) result.push(args[index]);
  }
  return result;
}

async function collectGlbs(target) {
  const resolved = path.resolve(projectRoot, target);
  const details = await stat(resolved);
  if (details.isFile()) return resolved.toLowerCase().endsWith(".glb") ? [resolved] : [];
  const files = [];
  for (const entry of await readdir(resolved, { withFileTypes: true })) {
    const child = path.join(resolved, entry.name);
    if (entry.isDirectory()) files.push(...await collectGlbs(child));
    else if (entry.isFile() && !entry.name.startsWith(".") && entry.name.toLowerCase().endsWith(".glb")) files.push(child);
  }
  return files;
}

function rgba8MipBytes(width, height, levels) {
  let bytes = 0;
  for (let level = 0; level < levels; level += 1) {
    bytes += Math.max(1, width >> level) * Math.max(1, height >> level) * 4;
  }
  return bytes;
}

function blockCompressedMipBytes(width, height, levels, bytesPerBlock) {
  let bytes = 0;
  for (let level = 0; level < levels; level += 1) {
    const levelWidth = Math.max(1, width >> level);
    const levelHeight = Math.max(1, height >> level);
    bytes += Math.ceil(levelWidth / 4) * Math.ceil(levelHeight / 4) * bytesPerBlock;
  }
  return bytes;
}

function fullMipLevelCount(width, height) {
  return Math.floor(Math.log2(Math.max(width, height))) + 1;
}

async function readAt(handle, length, position) {
  const buffer = Buffer.allocUnsafe(length);
  const { bytesRead } = await handle.read(buffer, 0, length, position);
  if (bytesRead !== length) throw new Error(`Unexpected EOF at byte ${position}.`);
  return buffer;
}

export async function auditGlbTextures(file) {
  const handle = await open(file, "r");
  try {
    const header = await readAt(handle, 20, 0);
    if (header.toString("utf8", 0, 4) !== "glTF") throw new Error(`${file} is not a GLB.`);
    const jsonLength = header.readUInt32LE(12);
    const json = JSON.parse((await readAt(handle, jsonLength, 20)).toString("utf8"));
    const binHeaderOffset = 20 + jsonLength;
    const binHeader = await readAt(handle, 8, binHeaderOffset);
    if (binHeader.toString("utf8", 4, 8) !== "BIN\0") throw new Error(`${file} has no binary chunk.`);
    const binOffset = binHeaderOffset + 8;

    let ktx2Textures = 0;
    let etc1sTextures = 0;
    let uastcTextures = 0;
    let fallbackTextures = 0;
    let originalEquivalentBytes = 0;
    let compressedGpuBytes = 0;
    for (const image of json.images ?? []) {
      if (image.mimeType !== "image/ktx2" || image.bufferView === undefined) {
        fallbackTextures += 1;
        if (image.bufferView !== undefined) {
          const view = json.bufferViews?.[image.bufferView];
          if (!view) throw new Error(`${file} image references missing buffer view ${image.bufferView}.`);
          const source = await readAt(handle, view.byteLength, binOffset + (view.byteOffset ?? 0));
          const metadata = await sharp(source).metadata();
          if (metadata.width && metadata.height) {
            const fallbackBytes = rgba8MipBytes(
              metadata.width,
              metadata.height,
              fullMipLevelCount(metadata.width, metadata.height),
            );
            originalEquivalentBytes += fallbackBytes;
            compressedGpuBytes += fallbackBytes;
          }
        }
        continue;
      }
      const view = json.bufferViews?.[image.bufferView];
      if (!view) throw new Error(`${file} image references missing buffer view ${image.bufferView}.`);
      const ktxHeader = await readAt(handle, 48, binOffset + (view.byteOffset ?? 0));
      if (!ktxHeader.subarray(0, 12).equals(KTX2_IDENTIFIER)) {
        throw new Error(`${file} declares image/ktx2 without a KTX2 identifier.`);
      }
      const width = ktxHeader.readUInt32LE(20);
      const height = ktxHeader.readUInt32LE(24);
      const levels = Math.max(1, ktxHeader.readUInt32LE(40));
      const supercompressionScheme = ktxHeader.readUInt32LE(44);
      const bytesPerBlock = supercompressionScheme === 1 ? 8 : 16;
      if (supercompressionScheme === 1) etc1sTextures += 1;
      else uastcTextures += 1;
      originalEquivalentBytes += rgba8MipBytes(width, height, levels);
      compressedGpuBytes += blockCompressedMipBytes(width, height, levels, bytesPerBlock);
      ktx2Textures += 1;
    }

    return {
      file,
      fileBytes: (await stat(file)).size,
      ktx2Textures,
      etc1sTextures,
      uastcTextures,
      fallbackTextures,
      originalEquivalentBytes,
      compressedGpuBytes,
    };
  } finally {
    await handle.close();
  }
}

async function main() {
  const args = process.argv.slice(2);
  const visualOnly = args.includes("--visual-only");
  const creaturesOnly = args.includes("--creatures-only");
  const requireAll = args.includes("--require-all");
  const reportPath = argumentValue(args, "--report", null);
  const targets = positionalArguments(args);
  if (!targets.length) targets.push("src/assets/game/dungeons", "public/assets/game/manastorm");
  const files = [...new Set((await Promise.all(targets.map(collectGlbs))).flat())]
    .filter((file) => !visualOnly || /(?:^|[-_])visual\.glb$/i.test(path.basename(file)))
    .filter((file) => !creaturesOnly || file.split(path.sep).includes("creatures"))
    .sort();

  const results = [];
  for (const file of files) results.push(await auditGlbTextures(file));
  const totals = results.reduce((sum, result) => ({
    files: sum.files + 1,
    fileBytes: sum.fileBytes + result.fileBytes,
    ktx2Textures: sum.ktx2Textures + result.ktx2Textures,
    etc1sTextures: sum.etc1sTextures + result.etc1sTextures,
    uastcTextures: sum.uastcTextures + result.uastcTextures,
    fallbackTextures: sum.fallbackTextures + result.fallbackTextures,
    originalEquivalentBytes: sum.originalEquivalentBytes + result.originalEquivalentBytes,
    compressedGpuBytes: sum.compressedGpuBytes + result.compressedGpuBytes,
  }), {
    files: 0,
    fileBytes: 0,
    ktx2Textures: 0,
    etc1sTextures: 0,
    uastcTextures: 0,
    fallbackTextures: 0,
    originalEquivalentBytes: 0,
    compressedGpuBytes: 0,
  });
  const reduction = totals.originalEquivalentBytes > 0
    ? 1 - totals.compressedGpuBytes / totals.originalEquivalentBytes
    : 0;
  console.log(
    `KTX2 audit: ${totals.files} files, ${totals.ktx2Textures} compressed textures `
    + `(${totals.etc1sTextures} ETC1S, ${totals.uastcTextures} UASTC), `
    + `${totals.fallbackTextures} fallback textures; estimated texture GPU memory `
    + `${(totals.originalEquivalentBytes / 1_048_576).toFixed(1)} -> `
    + `${(totals.compressedGpuBytes / 1_048_576).toFixed(1)} MiB (${(reduction * 100).toFixed(1)}% reduction).`,
  );

  if (reportPath) {
    const absoluteReportPath = path.resolve(projectRoot, reportPath);
    await mkdir(path.dirname(absoluteReportPath), { recursive: true });
    await writeFile(absoluteReportPath, `${JSON.stringify({
      schemaVersion: 1,
      generatedAt: new Date().toISOString(),
      totals,
      estimatedGpuReduction: reduction,
      files: results.map((result) => ({
        ...result,
        file: path.relative(projectRoot, result.file).replaceAll("\\", "/"),
      })),
    }, null, 2)}\n`, "utf8");
  }
  if (requireAll && totals.fallbackTextures > 0) process.exitCode = 1;
}

if (process.argv[1] && path.resolve(process.argv[1]) === fileURLToPath(import.meta.url)) {
  await main();
}