# Copyright 2021 Joe Drago. All rights reserved.
# SPDX-License-Identifier: BSD-2-Clause

# READ THIS WHOLE COMMENT FIRST, BEFORE RUNNING THIS SCRIPT:

# The goal of this script is to detect AVIFs containing multiple adjacent iref boxes and merge them,
# filling leftover space with a free box (to avoid ruining file offsets elsewhere in the file). The
# syntax is simple:

#     coffee irefmerge.coffee filename.avif

# This will look over the file's contents and if it detects multiple irefs, it will fix it in
# memory, make a adjacent backup of the file (filename.avif.irefmergeBackup), and then overwrite the
# original file with the fixed contents. Using -v on the commandline will enable Verbose mode, and
# using -n will disable the creation of backups (.irefmergeBackup files).

# This should be well-behaved on files created by old versions of avifenc, but **PLEASE** make
# backups of your images before running this script on them, **especially** if you plan to run with
# "-n". I do not advise running this script on AVIFs generated by anything other than avifenc.

# Possible responses for a file:
# * [NotAvif] This file isn't an AVIF.
# * [BadAvif] This file thinks it is an AVIF, but is missing important things.
# * [Skipped] This file is an AVIF, but didn't need any fixes.
# * [Success] This file is an AVIF, had to be fixed, and was fixed.
# * (the script crashes) I probably have a bug; let me know.

# Note on CoffeeScript:
# If you don't want to invoke coffeescript every time, you can compile it once with:
#     coffee -c -b irefmerge.coffee
# ... and run the adjacent irefmerge.js with node instead. "It's just JavaScript."

# -------------------------------------------------------------------------------------------------
# Syntax

syntax = ->
  console.log "Syntax: irefmerge [-v] [-n] file1 [file2 ...]"
  console.log "        -v : Verbose mode"
  console.log "        -n : No Backups (Don't generate adjacent .irefmergeBackup files when overwriting in-place)"

# -------------------------------------------------------------------------------------------------
# Constants and helpers

fs = require 'fs'

INDENT = "         "
VERBOSE = false

verboseLog = ->
  if VERBOSE
    console.log.apply(null, arguments)

fatalError = (reason) ->
  console.error "ERROR: #{reason}"
  process.exit(1)

# -------------------------------------------------------------------------------------------------
# Box

class Box
  constructor: (@filename, @type, @buffer, @start, @size) ->
    @offset = @start
    @bytesLeft = @size
    @version = 0
    @flags = 0
    @boxes = {} # child boxes

  nextBox: ->
    if @bytesLeft < 8
      return null
    boxSize = @buffer.readUInt32BE(@offset)
    boxType = @buffer.toString('utf8', @offset + 4, @offset + 8)
    if boxSize > @bytesLeft
      verboseLog("#{INDENT} * Truncated box of type #{boxType} (#{boxSize} bytes with only #{@bytesLeft} bytes left)")
      return null
    if boxSize < 8
      verboseLog("#{INDENT} * Bad box size of type #{boxType} (#{boxSize} bytes")
      return null
    newBox = new Box(@filename, boxType, @buffer, @offset + 8, boxSize - 8)
    @offset += boxSize
    @bytesLeft -= boxSize
    verboseLog "#{INDENT} * Discovered box type: #{newBox.type} offset: #{newBox.offset - 8} size: #{newBox.size + 8}"
    return newBox

  walkBoxes: ->
    while box = @nextBox()
      @boxes[box.type] = box
    return

  readFullBoxHeader: ->
    if @bytesLeft < 4
      fatalError("#{INDENT} * Truncated FullBox of type #{boxType} (only #{@bytesLeft} bytes left)")
    versionAndFlags = @buffer.readUInt32BE(@offset)
    @version = (versionAndFlags >> 24) & 0xFF
    @flags = versionAndFlags & 0xFFFFFF
    @offset += 4
    @bytesLeft -= 4
    return

  ftypHasBrand: (brand) ->
    if @type != 'ftyp'
      fatalError("#{INDENT} * Calling Box.ftypHasBrand() on a non-ftyp box")
    majorBrand = @buffer.toString('utf8', @offset, @offset + 4)
    compatibleBrands = []
    compatibleBrandCount = Math.floor((@size - 8) / 4)
    for i in [0...compatibleBrandCount]
      o = @offset + 8 + (i * 4)
      compatibleBrand = @buffer.toString('utf8', o, o + 4)
      compatibleBrands.push compatibleBrand

    verboseLog "#{INDENT}   * ftyp majorBrand: #{majorBrand} compatibleBrands: [#{compatibleBrands.join(', ')}]"

    if majorBrand == brand
      return true
    for compatibleBrand in compatibleBrands
      if compatibleBrand == brand
        return true
    return false

# -------------------------------------------------------------------------------------------------
# Main

irefMerge = (filename, makeBackups) ->
  if not fs.existsSync(filename)
    fatalError("File doesn't exist: #{filename}")
  try
    fileBuffer = fs.readFileSync(filename)
  catch e
    fatalError "Failed to read \"#{filename}\": #{e}"

  fileBox = new Box(filename, "<file>", fileBuffer, 0, fileBuffer.length)
  fileBox.walkBoxes()

  ftypBox = fileBox.boxes.ftyp
  if not ftypBox?
    return "NotAvif"
  if ftypBox.type != 'ftyp'
    return "NotAvif"
  if !ftypBox.ftypHasBrand('avif')
    return "NotAvif"

  metaBox = fileBox.boxes.meta
  if not metaBox?
    return "BadAvif"
  metaBox.readFullBoxHeader()

  merged = false
  irefs = []
  while box = metaBox.nextBox()
    if box.type == 'iref'
      irefs.push box

  # console.log irefs
  if irefs.length > 1
    verboseLog "#{INDENT} * Discovered multiple (#{irefs.length}) iref boxes, merging..."
    # merge irefs, and leave a free block in the dead space
    newTotalSize = 8 + 4 # the new single iref header's size + fullbox
    for iref in irefs
      newTotalSize += iref.size - 4
    fileBuffer.writeUInt32BE(newTotalSize, irefs[0].start - 8)

    writeOffset = irefs[0].start + 4 # skip past the fullbox's version[1]+flags[3]
    for iref in irefs
      fileBuffer.copy(fileBuffer, writeOffset, iref.start + 4, iref.start + iref.size)
      writeOffset += iref.size - 4
    freeBoxSize = (irefs.length - 1) * 12
    freeBox = Buffer.alloc(freeBoxSize)
    freeBox.fill(0)
    freeBox.writeUInt32BE(freeBoxSize)
    freeBox.write("free", 4)
    freeBox.copy(fileBuffer, writeOffset, 0, freeBoxSize)
    verboseLog "#{INDENT}   * Wrote a free chunk of size #{freeBoxSize} at offset #{writeOffset}"
    merged = true

  if merged
    if makeBackups
      backupFilename = filename + ".irefmergeBackup"
      fs.writeFileSync(backupFilename, fs.readFileSync(filename))
    fs.writeFileSync(filename, fileBuffer)
    return "Success"
  return "Skipped"

main = ->
  showSyntax = false
  makeBackups = true
  files = []

  for arg in process.argv.slice(2)
    switch arg
      when '-h', '--help'
        showSyntax = true
        break
      when '-n', '--no-backups'
        makeBackups = false
        break
      when '-v', '--verbose'
        VERBOSE = true
        break
      else
        files.push arg

  if showSyntax or files.length == 0
    return syntax()

  for filename in files
    verboseLog("[Reading] #{filename}")
    result = irefMerge(filename, makeBackups)
    console.log("[#{result}] #{filename}") # Always print this

  return 0

main()