#!/usr/bin/env ruby

if ARGV.size != 1
  $stderr.puts "Usage: extract_wb <file>"
  exit 2
end

# Presets we just skip
IGNORED_PRESETS=["Auto","Kelvin","Measured","AsShot","Kelvin"]
# Preset names we adjust
FL_PRESET_REPLACE={
  "Fluorescent" => "CoolWhiteFluorescent",
  "FluorescentP1" => "DayWhiteFluorescent",
  "FluorescentP2" => "DaylightFluorescent",
  "FluorescentM1" => "WarmWhiteFluorescent",
  "FluorescentD"  => "DaylightFluorescent",
  "FluorescentN"  => "NeutralFluorescent",
  "FluorescentW"  => "WhiteFluorescent",
}
PRESET_ORDER=["Unknown", "Daylight","Shade","Cloudy","Tungsten","Incandescent","Fluorescent",
              "WarmWhiteFluorescent","CoolWhiteFluorescent","DayWhiteFluorescent","DaylightFluorescent",
              "DaylightFluorescent", "NeutralFluorescent", "WhiteFluorescent", "Flash"]

PRESET_SORT_MAPPING={}
PRESET_ORDER.each_with_index {|name, index| PRESET_SORT_MAPPING[name]=index+1}

CAMERAS=File.expand_path("../src/external/rawspeed/data/cameras.xml", File.dirname(__FILE__))

# Create a map between EXIF maker/model and our cleaned up maker/model
require 'nokogiri'
exif_name_map = {}
File.open(CAMERAS) do |f|
  xml_doc  = Nokogiri::XML(f)
  xml_doc.css("Camera").each do |c|
    maker = exif_maker = c.attribute("make").value
    model = exif_model = c.attribute("model").value
    exif_id = [maker,model]
    if c.css("ID")[0]
      maker = c.css("ID")[0].attribute("make").value
      model = c.css("ID")[0].attribute("model").value
    end
    exif_name_map[exif_id] = [maker,model]
    c.css("Alias").each do |a|
      exif_model = a.content
      exif_id = [exif_maker, exif_model]
      exif_name_map[exif_id] = [maker, model]
    end
  end
end

red = green = blue = maker = model = preset = nil
listed_presets = []
fl_count = 0
IO.popen("exiftool \"-WB_RG*B*\" -WhiteBalance -Make -Model #{ARGV[0]}") do |io|
  while !io.eof?
    lineparts = io.readline.split(":")
    tag = lineparts[0].strip
    values = lineparts[1].strip.split(" ")

    if tag == "WB RGGB Levels"
      green = (values[1].to_f+values[2].to_f)/2.0
      red = values[0].to_f/green
      blue = values[3].to_f/green
      green = 1
    elsif tag.split.include? "Make"
      maker = lineparts[1].strip
    elsif tag.split.include? "Model"
      model = lineparts[1].strip
    elsif tag == "White Balance"
      preset = values[0]
    elsif ["WB RGB Levels", "WB RGGB Levels"].include? tag.split[0..2].join(" ")
      p = tag.split[3..-1].join

      r=g=b=0
      if values.size == 4
        g = (values[1].to_f+values[2].to_f)/2.0
        r = values[0].to_f/g
        b = values[3].to_f/g
        g = 1
      elsif values.size == 3
        g = values[1].to_f
        r = values[0].to_f/g
        b = values[2].to_f/g
        g = 1
      else
        $stderr.puts "Found tag '#{p}' with #{values.size} values instead of 3 or 4"
      end
      fl_count += 1 if p[0..."Fluorescent".size] == "Fluorescent"
      p = "Unknown" if !p || p==""
      listed_presets << [p,r,g,b]
    end
  end
end

if fl_count > 1
  listed_presets = listed_presets.map do |preset, red, green, blue|
    preset = FL_PRESET_REPLACE[preset] if FL_PRESET_REPLACE[preset]
    [preset, red, green, blue]
  end
end

# Adjust the maker/model we found with the map we generated before
if exif_name_map[[maker,model]]
  maker, model = exif_name_map[[maker,model]]
else
  $stderr.puts "WARNING: Couldn't find model in cameras.xml ('#{maker}', '#{model}')"
end

def print_preset(maker, model, preset, red, green, blue)
  if IGNORED_PRESETS.include? preset
    $stderr.puts "Ignoring preset '#{preset}'"
    return
  end

  preset = '"'+preset+'"' if preset[-1] == "K"

  puts "  { \"#{maker}\", \"#{model}\", #{preset}, 0, { #{red}, #{green}, #{blue}, 0 } },"
end

# Print out the WB value that was used in the file
preset = "\"#{ARGV[0]}\"" if !preset
if red && green && blue
  print_preset(maker, model, preset, red, green, blue)
  puts
else
  $stderr.puts "ERROR: Couldn't figure out preset values, sorry"
end

def convert_preset_to_sort(preset)
  if IGNORED_PRESETS.include? preset
    return 0
  elsif PRESET_SORT_MAPPING[preset]
    return PRESET_SORT_MAPPING[preset]
  elsif preset[-1] == "K"
    return preset[0..-2].to_i
  else
    $stderr.puts "WARNING: no defined sort order for '#{preset}'"
    return 0
  end
end

listed_presets.sort! do |a, b|
  convert_preset_to_sort(a[0]) <=> convert_preset_to_sort(b[0])
end

# Print out the presets if they exit
listed_presets.each do |preset, red, green, blue|
  print_preset(maker, model, preset, red, green, blue)
end
