#!/usr/bin/env ruby

require 'pp'
require 'digest/md5'

module Obj
  class Header
    SIZE = 5*4 # 5x 32-bit little-endian integers

    attr_reader(
      :num_sprites, # Number of entries in the directory. Always dir_size/8?
      :dir_offset,  # Offset of the sprite directory. Also the size of the main header
      :dir_size,    # Number of bytes allocated to the sprite directory. Always num_sprites*8?
      :data_offset, # Number of bytes allocated to the sprites themselves
      :data_size    # Total size of the sprite data
      # JUNGTIL.obj has 0x0000 0x0000 0x0000 following. Probably just padding to
      # allow the sprite directory to start at a word boundary.
    )

    def self.parse(data)
      hdr = new(*data[0..SIZE - 1].unpack("VVVVV"))
      # pp hdr
      hdr.validate!(data.bytes.size)
      hdr
    end

    def initialize(*entries)
      @num_sprites, @dir_offset, @dir_size, @data_offset, @data_size = *entries
    end

    def to_data
      [
        @num_sprites,
        @dir_offset,
        @dir_size,
        @data_offset,
        @data_size,
      ].pack("VVVVV")
    end

    def validate!(overall_size)
      raise "Directory overlaps EOF" if overall_size < dir_size + dir_offset
      raise "Data overlaps EOF" if overall_size < data_offset + data_size
      raise "Bad dir_size or num_sprites" if num_sprites * DirEntry::SIZE != dir_size
      raise "Directory overlaps data" if data_range.cover?(dir_offset+1)
      raise "Data overlaps directory" if dir_range.cover?(data_offset+1)

      true
    end

    def dir_range
      dir_offset...(dir_offset+dir_size)
    end

    def data_range
      data_offset...(data_offset+data_size)
    end
  end

  class DirEntry
    SIZE = 8

    attr_reader :rel_offset # Relative to the main header's data_offset
    attr_reader :sprite_size

    def self.parse(rel_data)
      new(*rel_data.unpack("VV"))
    end

    def initialize(rel_offset, sprite_size)
      @rel_offset = rel_offset
      @sprite_size = sprite_size
    end

    def to_data
      [rel_offset, sprite_size].pack("VV")
    end

    def sprite_range
      rel_offset...(rel_offset+sprite_size)
    end
  end

  class SpriteDir
    attr_reader :entries

    def self.parse(rel_data)
      count = rel_data.bytes.size
      num_entries = count / DirEntry::SIZE
      trailing = count%DirEntry::SIZE

      raise "SpriteDir block has #{trailing} trailing bytes" unless trailing == 0

      entries = 0.upto(num_entries-1).map do |n|
        rel_offset = n * DirEntry::SIZE
        DirEntry.parse(rel_data.byteslice(rel_offset, DirEntry::SIZE))
      end

      # pp entries

      new(entries)
    end
    
    def initialize(entries)
      @entries = entries
    end

    def to_data
      entries.map(&:to_data).join("")
    end

    # Convert the directory into an Array of sprites
    def realize(rel_data)
      entries.map { |entry| Sprite.parse(rel_data[entry.sprite_range]) }
    end
  end

  class SpriteHeader
    SIZE = 24        # Seems to be, anyway. Based on the size here vs. the size in the dir

    attr_reader(
      :unknown0,        # Possibly magic data? It's the same for every sprite in jungtil.obj
      :width, # Low nibble comes to 63 in jungtil.obj which would work for a 64x64 tile
      :height,
      :unknown8,
      :size,            # Number of bytes of pixel data following this header
      :unknown16,
      :unknown20,
    )

    def self.parse(rel_data)
      new(*rel_data[0...SIZE].unpack("VvvVVVV"))
    end

    def initialize(*args)
      @unknown0,
        @width,
        @height,
        @unknown8,
        @size,
        @unknown16,
        @unknown20 = *args
    end

    def to_data
      [
        @unknown0,
        @width,
        @height,
        @unknown8,
        @size,
        @unknown16,
        @unknown20
      ].pack("VvvVVVV")
    end

    def pixel_range
      SIZE...(SIZE+size)
    end
  end

  class Sprite
    attr_reader :header, :records, :raw

    def self.parse(rel_data)
      hdr = SpriteHeader.parse(rel_data)
      sprite_pixels = rel_data[hdr.pixel_range]

      records = sprite_pixels.split("\x00")
      records.map! { |record| record += "\x00" }

      new(hdr, records, rel_data)
    end

    def initialize(header, records, raw = nil)
      @header = header
      @records = records
      @raw = raw
    end

    # raw is optional, so don't use it here
    def to_data
      header.to_data + records.join("")
    end

    def total_size
      SpriteHeader::SIZE + records.join("").size
    end
  end

  class Parsed
    attr_reader :header
    attr_reader :directory
    attr_reader :sprites
    attr_accessor :filename

    def initialize(header, directory, sprites)
      @header = header
      @directory = directory
      @sprites = sprites
    end

    def to_data
      dir_padding = "\x00"*(header.dir_offset - Header::SIZE)
      data_padding = "" # for now, assume all is well

      [
        header.to_data,
        dir_padding,
        directory.to_data,
        data_padding,
        sprites.map { |sprite| sprite.to_data },
      ].join("")
    end
  end

  def self.parse(data)
    hdr = Header.parse(data)
    dir = SpriteDir.parse(data[hdr.dir_range])
    sprites = dir.realize(data[hdr.data_range])

    Parsed.new(hdr, dir, sprites)
  end
end

def hex(num, leading=0)
  return " "*(leading) if num.nil?

  "%0#{leading}x"%num
end

VISIBLE_CHARS = ("0".ord)..("z".ord)
def text(num)
  return " " if num.nil?
  if VISIBLE_CHARS.cover?(num)
    num.chr
  else
    " "
  end
end

def ascii(num)
  return "   " if num.nil?

  "%3d"%num
end

def header!(blocksize=8)
  nums = (0...blocksize).to_a
  hdr = "        #{nums.join("  ")}       #{nums.join("")}"

  puts "-" * (hdr.size + 1)
  puts hdr
  puts "-" * (hdr.size + 1)
end

def display(data, blocksize=8, skip=0, header: false)
  bytes = data.bytes[skip..-1]
  nrows = (bytes.count / blocksize)

  0.upto(nrows) do |i|
    header!(blocksize) if header && (i%16==0 || i == skip)
    block = bytes[(i*blocksize)...(i*blocksize+blocksize)]

    block.concat([nil]*(blocksize-block.size)) if block.size < blocksize

    out = [
      "0x#{hex(i*blocksize, 4)}",
      block.map { |b| hex(b, 2) },                          # hex
 #     "    | " + block.map { |b| text(b) }.join("") + " |",   # ascii
#      block.map { |b| ascii(b) } ,# decimal bytes
      "",# decimal 2-bytes
     # decimal 4-bytes
    ]

    puts out.flatten.join(' ')
    # puts Digest::MD5.hexdigest(data[skip..-1])
  end
end

def load_obj(filename)
  Obj.parse(File.read(filename).force_encoding("BINARY"))
end

def dump(filename, spriteno = -1)
  obj = load_obj(filename)
#  pp obj

  if spriteno == -1
    puts "Dumping all sprites for #{filename}"
    obj.sprites.each_with_index { |sprite, i| puts "Sprite #{i}\n" ; display(sprite.raw, 8, header: true) }
  else
    puts "Dumping sprite #{spriteno} for #{filename}"
    display(obj.sprites[spriteno])
  end
end

# The per-sprite data in jungtil.obj is too small to represent a 64x64 block,
# even at 1bpp, so try some decompression algorithms to see what comes out.
#
# Tried so far, with no success:
#   * DEFLATE / ZLIB
#   * LZO
#   * LZOP
#   * RLE (maybe?)
#
# Maybe try:
#   * RLE8: https://www.fileformat.info/format/bmp/corion-rle8.htm

def decompress(filename)
  puts "\nAttempting custom decompression of #{filename}..."

  obj = load_obj(filename)

  obj.sprites.each_with_index do |sprite, i|
    hdr = sprite.header
    print "Sprite %02d... x=#{hdr.width} y=#{hdr.height}"%i

    decompressed = []
    data = sprite.raw[24..-1].bytes
    hdr = sprite.header

    loop do
      break if data.empty?

      right = data.index(0)

      if right.nil?
        print "error! end of chunk not found"
        break
      end

      rec = data.shift(right)
      _ = data.shift(1) # drop the record separator

      decompressed << rec
    end
    puts ": #{decompressed.size} records. Sprite pixels: #{hdr.width*hdr.height}"

    puts "WARNING: #{data.size} bytes left over" if data.size > 0

    header!
    decompressed.each_with_index do |line, y|
      x = line.size
      padding = (sprite.header.width - x) / 2
      display((["\x00"]*padding + line.map(&:chr)).join(""), line.size+padding.size)
    end
  end
end

# Print the first 128 bytes of sprites in a friendly format. Permits into-obj comparisons
def compare(filenames)
  objs = filenames.map { |filename| load_obj(filename).tap { |obj| obj.filename = filename } }

  objs.each do |obj|

  end

  sprites = 
    if spriteno == -1
    else
    

    end
end
require 'set'
def correlate(filenames)
  objs = filenames.map { |f| load_obj(f) }

  results = Hash.new { |h,k| h[k] = Set.new }


  objs.each do |obj|
    obj.sprites.each do |spr|
      results[spr.header.unknown0 & 0xffff] << spr.header.width
    end
  end

  puts "Unique widths for u0,1"
  pp results
end

def directory(filename, num)
  data = File.read(filename).force_encoding("BINARY")

  hdr = Obj::Header.parse(data)
  dir = Obj::SpriteDir.parse(data[hdr.dir_range])
  entry = dir.entries[num]

  puts "Sprite directory starts at 0x#{hdr.dir_offset.to_s(16)}"
  puts "Directory entry for sprite #{num} is at 0x#{(hdr.dir_offset + (Obj::DirEntry::SIZE * num)).to_s(16)}"
  puts "Sprite #{num} is at 0x#{(hdr.data_offset + entry.rel_offset).to_s(16)} and is #{entry.sprite_size} bytes"
end

def sprites(filename)
  obj = load_obj(filename)

  puts filename + ":"
  obj.sprites.each_with_index do |spr, i|
    hdr = spr.header
    px = hdr.width * hdr.height
    pp hdr
    pp spr

    data = spr.to_data

    puts "%s %03d: %02x %02x %02x %02x %3d %3d %5d %5d  %.2f     %02x %02x ... %02x"%[
      filename,
      i,
      hdr.unknown0 & 0x000000ff,
      (hdr.unknown0 & 0x0000ff00) >> 8,
      (hdr.unknown0 & 0x00ff0000) >> 16,
      (hdr.unknown0 & 0xff000000) >> 24,
      hdr.width,
      hdr.height,
      px,
      hdr.size,
      (hdr.size*8) / px.to_f,
      data.bytes[0],
      data.bytes[1],
      data.bytes[-1],
    ]
  end
end

def sprite(filename, i)
  obj = load_obj(filename)

  puts filename + ":"
  spr = obj.sprites[i]
  raise "No sprite #{i}" unless spr

  # Show the header
  display(spr.raw[0...24], 4, header: true)

  spr.records.each_with_index do |record, i|
    str =
      if record.size > 40
        record[0...20].bytes.map { |b| hex(b, 2) }.join(" ") +
        " ... " +
        record[-20..-1].bytes.map { |b| hex(b, 2) }.join(" ")
      else
        record.bytes.map { |b| hex(b, 2) }.join(" ")
      end

    puts "%4d: %4d bytes: %s"%[i,record.size, str]
  end
end

# Build a test sprite to investigate the color palette
def build(filename)
  sprite_records = [
    # 63 in each sprite except the last, which has just 3
      1.upto(63).map { |b| "\x01#{b.chr}\x00" },
     64.upto(126).map { |b| "\x01#{b.chr}\x00" },
    127.upto(189).map { |b| "\x01#{b.chr}\x00" },
    190.upto(252).map { |b| "\x01#{b.chr}\x00" },
    253.upto(255).map { |b| "\x01#{b.chr}\x00" },
  ]

  sprites = sprite_records.map do |records|
    data = records.join("")
    header = Obj::SpriteHeader.new(
       0, # Stolen from blank.obj
       1, # width of 1
      records.size, # height = number of records
       0, # padding?
       data.size,
       0, # padding?
       0, # padding?
    )

    Obj::Sprite.new(
      header,
      records,
      header.to_data + data,
    )
  end

  dir_block_offset = 32 # hardcoded
  dir_block_size = 8 *sprites.size
  data_block_offset = dir_block_offset + dir_block_size
  data_block_size = sprites.inject(0) { |x, spr| x + spr.total_size }

  header = Obj::Header.new(
      sprites.size,
      dir_block_offset,
      dir_block_size,
      data_block_offset,
      data_block_size,
  )

  offset = 0
  directory = Obj::SpriteDir.new(
    sprites.map { |sprite|
      puts "offset #{offset} => #{sprite.total_size}"
      entry = Obj::DirEntry.new(offset, sprite.total_size)
      offset += sprite.total_size
      entry
    }
  )

  pp directory

  built = Obj::Parsed.new(
    header,
    directory,
    sprites
  )

  File.open(filename, "w") { |f| f.write(built.to_data) }
end

def unknown16(filenames)
  objs = filenames.map { |f| load_obj(f) }
  results = Set.new

  objs.each do |obj|
    obj.sprites.each do |spr|
      results << spr.header.unknown16
    end
  end

  puts "Unique widths for u16,4"
  pp results
end

case command = ARGV.shift
when "directory" then
  directory(ARGV[0], ARGV[1].to_i)
when "unknown16" then
  unknown16(ARGV)
when "sprites" then
  ARGV.each { |filename| sprites(filename) } 
when "sprite" then
  sprite(ARGV[0], ARGV[1].to_i)
when "dump" then
  ARGV.each { |filename| dump(filename) } 
when "compare" then
  compare(ARGV)
when "decompress" then
  ARGV.each { |filename| decompress(filename) }
when "correlate" then
  correlate(ARGV)
when "build" then
  build(ARGV[0])
else
  puts "Unrecognized command #{command}"
  exit(1)
end