#!/usr/bin/env node

import { createHash } from "node:crypto";
import { createReadStream } from "node:fs";
import {
  lstat,
  mkdir,
  readdir,
  readFile,
  realpath,
  rename,
  rm,
  stat,
  writeFile,
} from "node:fs/promises";
import path from "node:path";
import process from "node:process";
import { isMap, parseDocument } from "yaml";

type CatalogFile = {
  path?: unknown;
  size?: unknown;
  sha256?: unknown;
};

type ProgramIndex = {
  type?: unknown;
  versions?: Array<{ files?: CatalogFile[] }>;
  additional_files?: CatalogFile[];
};

type Options = {
  check: boolean;
  targets: string[];
};

type FileReference = {
  data: CatalogFile;
  yamlPath: Array<string | number>;
};

type ResolvedCatalogFile = {
  path: string;
  size: number;
};

const ignoredDirectories = new Set([".git", "node_modules", "files", "img"]);

function usage(): void {
  console.log(`Usage: pnpm fill-metadata [--check] [path ...]

With no paths, the script scans the current directory. A path may point to an
index.yaml file or a directory containing index.yaml files.

  --check  verify checksums and file sizes without modifying YAML
  --help   show this help message`);
}

function parseArguments(args: string[]): Options {
  const targets: string[] = [];
  let check = false;

  for (const argument of args) {
    if (argument === "--check") {
      check = true;
    } else if (argument === "--help" || argument === "-h") {
      usage();
      process.exit(0);
    } else if (argument.startsWith("-")) {
      throw new Error(`Unknown option: ${argument}`);
    } else {
      targets.push(argument);
    }
  }

  return { check, targets: targets.length > 0 ? targets : ["."] };
}

async function findIndexes(target: string, result: Set<string>): Promise<void> {
  const absoluteTarget = path.resolve(target);
  const targetStat = await lstat(absoluteTarget);

  if (targetStat.isFile()) {
    if (path.basename(absoluteTarget) !== "index.yaml") {
      throw new Error(`Expected an index.yaml file: ${target}`);
    }
    result.add(absoluteTarget);
    return;
  }

  if (!targetStat.isDirectory()) {
    throw new Error(`Path is neither a file nor a directory: ${target}`);
  }

  const entries = await readdir(absoluteTarget, { withFileTypes: true });
  for (const entry of entries) {
    const entryPath = path.join(absoluteTarget, entry.name);
    if (entry.isFile() && entry.name === "index.yaml") {
      result.add(entryPath);
    } else if (entry.isDirectory() && !ignoredDirectories.has(entry.name)) {
      await findIndexes(entryPath, result);
    }
  }
}

function collectFileReferences(data: ProgramIndex): FileReference[] {
  const references: FileReference[] = [];

  for (const [versionIndex, version] of (data.versions ?? []).entries()) {
    for (const [fileIndex, file] of (version.files ?? []).entries()) {
      references.push({
        data: file,
        yamlPath: ["versions", versionIndex, "files", fileIndex],
      });
    }
  }

  for (const [fileIndex, file] of (data.additional_files ?? []).entries()) {
    references.push({
      data: file,
      yamlPath: ["additional_files", fileIndex],
    });
  }

  return references;
}

async function resolveCatalogFile(
  indexPath: string,
  value: unknown,
): Promise<ResolvedCatalogFile> {
  if (typeof value !== "string" || value.length === 0) {
    throw new Error("the path field is missing or is not a string");
  }
  if (path.isAbsolute(value) || value.includes("\\")) {
    throw new Error(`invalid path: ${value}`);
  }

  const parts = value.split("/");
  if (parts[0] !== "files" || parts.some((part) => part === ".." || part === "")) {
    throw new Error(`path must be inside files/: ${value}`);
  }

  const programDirectory = path.dirname(indexPath);
  const filesDirectory = path.resolve(programDirectory, "files");
  const filePath = path.resolve(programDirectory, ...parts);
  const relativePath = path.relative(filesDirectory, filePath);
  if (relativePath.startsWith("..") || path.isAbsolute(relativePath)) {
    throw new Error(`path escapes files/: ${value}`);
  }

  const [realFilesDirectory, realFilePath] = await Promise.all([
    realpath(filesDirectory),
    realpath(filePath),
  ]);
  const realRelativePath = path.relative(realFilesDirectory, realFilePath);
  if (realRelativePath.startsWith("..") || path.isAbsolute(realRelativePath)) {
    throw new Error(`symbolic link escapes files/: ${value}`);
  }
  const fileStat = await stat(realFilePath);
  if (!fileStat.isFile()) {
    throw new Error(`path does not point to a regular file: ${value}`);
  }

  return { path: realFilePath, size: fileStat.size };
}

async function calculateSha256(filePath: string): Promise<string> {
  const hash = createHash("sha256");
  for await (const chunk of createReadStream(filePath)) {
    hash.update(chunk);
  }
  return hash.digest("hex");
}

async function writeAtomically(filePath: string, contents: string): Promise<void> {
  const directory = path.dirname(filePath);
  const temporaryPath = path.join(
    directory,
    `.${path.basename(filePath)}.${process.pid}.tmp`,
  );

  await mkdir(directory, { recursive: true });
  try {
    await writeFile(temporaryPath, contents, "utf8");
    await rename(temporaryPath, filePath);
  } finally {
    await rm(temporaryPath, { force: true });
  }
}

function setFileSize(
  document: ReturnType<typeof parseDocument>,
  yamlPath: Array<string | number>,
  size: number,
): void {
  document.setIn([...yamlPath, "size"], size);

  const fileMap = document.getIn(yamlPath, true);
  if (!isMap(fileMap)) {
    return;
  }

  const keyName = (pair: (typeof fileMap.items)[number]): string => String(pair.key);
  const sizeIndex = fileMap.items.findIndex((pair) => keyName(pair) === "size");
  const sha256Index = fileMap.items.findIndex((pair) => keyName(pair) === "sha256");
  if (sizeIndex < 0 || sha256Index < 0 || sizeIndex < sha256Index) {
    return;
  }

  const [sizePair] = fileMap.items.splice(sizeIndex, 1);
  fileMap.items.splice(sha256Index, 0, sizePair);
}

async function processIndex(indexPath: string, check: boolean): Promise<number> {
  const source = await readFile(indexPath, "utf8");
  const document = parseDocument(source, { prettyErrors: true });
  if (document.errors.length > 0) {
    throw new Error(document.errors.map((error) => error.message).join("; "));
  }

  const data = document.toJS() as ProgramIndex;
  if (data.type !== "program") {
    return 0;
  }

  let changed = 0;
  const errors: string[] = [];
  const references = collectFileReferences(data);

  const results = await Promise.all(
    references.map(async (reference) => {
      try {
        const file = await resolveCatalogFile(indexPath, reference.data.path);
        return {
          reference,
          digest: await calculateSha256(file.path),
          size: file.size,
        };
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        errors.push(`${String(reference.data.path)}: ${message}`);
        return undefined;
      }
    }),
  );

  if (errors.length > 0) {
    throw new Error(errors.join("; "));
  }

  for (const result of results) {
    if (result === undefined) {
      continue;
    }

    if (result.reference.data.size !== result.size) {
      changed += 1;
      if (!check) {
        setFileSize(document, result.reference.yamlPath, result.size);
      }
    }
    if (result.reference.data.sha256 !== result.digest) {
      changed += 1;
      if (!check) {
        document.setIn([...result.reference.yamlPath, "sha256"], result.digest);
      }
    }
  }

  const displayPath = path.relative(process.cwd(), indexPath) || indexPath;
  if (changed > 0 && check) {
    console.error(`${displayPath}: metadata update required (${changed} fields)`);
  } else if (changed > 0) {
    await writeAtomically(indexPath, document.toString({ lineWidth: 0 }));
    console.log(`${displayPath}: updated file metadata (${changed} fields)`);
  }

  return changed;
}

async function main(): Promise<void> {
  const options = parseArguments(process.argv.slice(2));
  const indexes = new Set<string>();

  for (const target of options.targets) {
    await findIndexes(target, indexes);
  }

  let changed = 0;
  let failed = 0;
  for (const indexPath of [...indexes].sort()) {
    try {
      changed += await processIndex(indexPath, options.check);
    } catch (error) {
      failed += 1;
      const displayPath = path.relative(process.cwd(), indexPath) || indexPath;
      const message = error instanceof Error ? error.message : String(error);
      console.error(`${displayPath}: ${message}`);
    }
  }

  if (failed > 0) {
    console.error(`Errors: ${failed}`);
    process.exitCode = 1;
  } else if (options.check && changed > 0) {
    console.error(`File metadata mismatches: ${changed}`);
    process.exitCode = 1;
  } else {
    console.log(
      changed > 0
        ? `Done. Updated fields: ${changed}`
        : `Done. All SHA-256 checksums and file sizes are up to date. Checked index.yaml files: ${indexes.size}`,
    );
  }
}

await main();