#!/usr/bin/env ruby
#
# jpeg_inject_default_huffman.rb - injects a default Huffman table
# into a JPEG if that JPEG is detected to be lacking the table.
#
# Takes JPEG file on commandline, and overwrite that file with the
# updated header including the Huffman table.  Probably a few bugs;
# use at your own risk.
#
# Copyright (c) 2008 Patrick Mahoney
# All rights reserved.
#
# This code was written by Patrick Mahoney (pat@polycrystal.org) as
# part of Google Summer of Code 2008.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of the author nor the names of any contributors
#    may be used to endorse or promote products derived from this
#    software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
# ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
# TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.


module JPEG

  # See JPEG spec ISO/IEC 10918-1:1993(E) from http://www.wotsit.org/

  START_OF_IMAGE = 0xffd8
  START_OF_FRAME = 0xffc0
  DEFINE_HUFFMAN = 0xffc4
  START_OF_SCAN = 0xffda

  # From ISO/IEC 10918-1:1993(E) K.3.3.  This is the default Huffman
  # table used by USB Video Class devices using the MJPEG format but
  # lacking a Huffman table.
  HUFFMAN_lumDC_LENGTHS = [0, 1, 5, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0]
  HUFFMAN_lumDC_VALUES = [0x00, 0x01, 0x02, 0x03,
                          0x04, 0x05, 0x06, 0x07,
                          0x08, 0x09, 0x0A, 0x0B]
  HUFFMAN_lumDC = [HUFFMAN_lumDC_LENGTHS, HUFFMAN_lumDC_VALUES].flatten
  HUFFMAN_chromDC_LENGTHS = [0, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]
  HUFFMAN_chromDC_VALUES = [0x00, 0x01, 0x02, 0x03,
                            0x04, 0x05, 0x06, 0x07,
                            0x08, 0x09, 0x0A, 0x0B]
  HUFFMAN_chromDC = [HUFFMAN_chromDC_LENGTHS, HUFFMAN_chromDC_VALUES].flatten

  HUFFMAN_lumAC_LENGTHS = [0x00, 0x02, 0x01, 0x03, 0x03, 0x02, 0x04, 0x03,
                           0x05, 0x05, 0x04, 0x04, 0x00, 0x00, 0x01, 0x7D]
  HUFFMAN_lumAC_VALUES =
    [0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12, 0x21, 0x31, 0x41, 0x06, 0x13, 0x51, 0x61, 0x07,
     0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xa1, 0x08, 0x23, 0x42, 0xb1, 0xc1, 0x15, 0x52, 0xd1, 0xf0,
     0x24, 0x33, 0x62, 0x72, 0x82, 0x09, 0x0A, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x25, 0x26, 0x27, 0x28,
     0x29, 0x2A, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49,
     0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69,
     0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89,
     0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7,
     0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5,
     0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA, 0xE1, 0xE2,
     0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8,
     0xF9, 0xFA]

  HUFFMAN_chromAC_LENGTHS = [0x00, 0x02, 0x01, 0x02, 0x04, 0x04, 0x03, 0x04,
                             0x07, 0x05, 0x04, 0x04, 0x00, 0x01, 0x02, 0x77]
  HUFFMAN_chromAC_VALUES =
    [0x00, 0x01, 0x02, 0x03, 0x11, 0x04, 0x05, 0x21, 0x31, 0x06, 0x12, 0x41, 0x51, 0x07, 0x61, 0x71,
     0x13, 0x22, 0x32, 0x81, 0x08, 0x14, 0x42, 0x91, 0xA1, 0xB1, 0xC1, 0x09, 0x23, 0x33, 0x52, 0xF0,
     0x15, 0x62, 0x72, 0xD1, 0x0A, 0x16, 0x24, 0x34, 0xE1, 0x25, 0xF1, 0x17, 0x18, 0x19, 0x1A, 0x26,
     0x28, 0x29, 0x2A, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x27,
     0x49, 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68,
     0x69, 0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
     0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5,
     0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3,
     0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA,
     0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8,
     0xF9, 0xFA]

  def append_default_huffman(str)
    huff_tables = Array.new
    # DC, lum
    huff_tables << [0x00, [HUFFMAN_lumDC_LENGTHS, HUFFMAN_lumDC_VALUES].flatten]
    # AC, lum
    huff_tables << [0x10, [HUFFMAN_lumAC_LENGTHS, HUFFMAN_lumAC_VALUES].flatten]
    # DC, chrom
    huff_tables << [0x01, [HUFFMAN_chromDC_LENGTHS, HUFFMAN_chromDC_VALUES].flatten]
    # AC, chrom
    huff_tables << [0x11, [HUFFMAN_chromAC_LENGTHS, HUFFMAN_chromAC_VALUES].flatten]

    huff_tables.each do |table|
      id = table[0]
      data = table[1]
      str.append_word(DEFINE_HUFFMAN)
      str.append_word(2 + 1 + data.size)
      str << id
      data.each {|b| str << b}
    end
  end
end

class String
  def append_word(word)
    self << (word >> 8)
    self << (word & 0xff)
  end
end

class IO
  def each_word
    first_in_pair = true
    high = 0
    last = nil
    each_byte do |b|
      if first_in_pair
        high = b
        last = b
        first_in_pair = false
      else
        yield((high << 8) | b)
        last = nil
        first_in_pair = true
      end
    end

    last  # in case there is an odd final byte
  end
end

include JPEG

UNKNOWN = 0
HEADER = 1
IMAGE = 3

state = UNKNOWN
out = ""

if ARGV.size != 1
  abort "usage: #{$0} file.jpg"
end

file = ARGV[0]

File.open(file) do |f|
  last_byte = f.each_word do |word|
    case state
    when UNKNOWN
      case word
      when START_OF_IMAGE
        state = HEADER
      else
        break
      end
    when HEADER
      case word
      when DEFINE_HUFFMAN
        abort "Huffman table already defined %x"
      when START_OF_SCAN
        append_default_huffman(out)
        state = IMAGE
      end
    end

    if state == UNKNOWN
      abort "Failed to detect JPEG"
    else
      out.append_word(word)
    end
  end

  out << last_byte if last_byte
end

File.open(file, "w") do |f|
  f.write out
  print "#{$0}: wrote JPEG with Huffman table to #{file}\n"
end