#!/usr/bin/ruby -w

# Parse a StepMania Stats.xml file and generate some personal high
# score tables...

require "rexml/document"

$convert_binary = "/usr/bin/convert"
$mkdir_binary = "/bin/mkdir"
$md5sum_binary = "/usr/bin/md5sum"

[ $convert_binary, $mkdir_binary, $md5sum_binary ].each do |binary|
  unless FileTest.executable? binary
    STDERR.puts "Couldn't find the binary: " + binary
    STDERR.puts "  (Edit the script to fix this...)"
    exit( -1 )
  end
end

def usage
  STDERR.print <<EOUSAGE
Usage:
    ./parse-stepmania-stats [OPTIONS] <stats-xml-filename> <output-directory>

  -h, --help           Display this message and exit.

  -m <feet>, --minimum-steps=<feet>
                       Only display high scores for songs of
                       difficulty of at least <feet>.

  -M <feet>, --maximum-steps=<feet>
                       Only display high scores for songs of
                       difficulty of no more than <feet> feet.

  -d, --include-doubles
                       Display scores for doubles as well as singles
                       steps types.

  -v, --verbose        Show extra information about the parsing as
                       we\'re going along.

EOUSAGE
end

require "getoptlong"

options = GetoptLong.new(

  [ "--help", "-h", GetoptLong::NO_ARGUMENT ],
  [ "--minimum-steps", "-m", GetoptLong::REQUIRED_ARGUMENT ],
  [ "--maximum-steps", "-M", GetoptLong::REQUIRED_ARGUMENT ],
  [ "--include-doubles", "-d", GetoptLong::NO_ARGUMENT ],
  [ "--verbose", "-v", GetoptLong::NO_ARGUMENT ]

)

$minimum_steps = 8
$maximum_steps = 10
$include_doubles = false

$verbose = false

begin

  options.each do |opt,arg|

    case opt
    when "--help"
      puts "got help..."
      usage
      exit( 0 )
    when "--minimum-steps"
      $minimum_steps = Integer( arg )
    when "--maximum-steps"
      $maximum_steps = Integer( arg )
    when "--include-doubles"
      $include_doubles = true
    when "--verbose"
      $verbose = true
    end

  end

rescue

  puts $!
  usage
  exit( -1 )

end

if $minimum_steps <= 1
  STDERR.puts "--minimum-steps must be at least 1"
  usage
  exit( -1 )
end

if $maximum_steps <= 1
  STDERR.puts "--maximum-steps must be at least 1"
  usage
  exit( -1 )
end

if $minimum_steps > $maximum_steps
  STDERR.puts "--minimum-steps must be less than or equal to --maximum_steps"
  usage
  exit( -1 )
end

unless ARGV.length == 2
  usage
  exit( -1 )
end

file_name = ARGV[0]
output_directory_name = ARGV[1]
output_index_name = output_directory_name + "/index.html.include"

unless system( $mkdir_binary, "-p", output_directory_name )
  STDERR.puts "Couldn't create output directory: " + output_directory_name
  exit( -1 )
end

full_file_name = File.expand_path( file_name )

$stepmania_root_directory_name = nil

if full_file_name =~ /^(.*)Data\/MachineProfile\/Stats.xml$/i
  $stepmania_root_directory_name = $1
else
  STDERR.puts file_name + " doesn't seem to be a Stats.xml " +
    " file from a StepMania installation."
  exit( -1 )
end

# A couple of helper functions for finding the MD5 sums of the MP3 and
# DWI files.

def safe_backticks( *command )

  result = ""

  Kernel.open( "|-", "r" ) do |f|
    if f
      f.each_line do |line|
        result += line
      end
    else
      begin
        exec( *command )
      rescue
        raise "Couldn't exec #{command}: #{$!}\n"
      end
    end
  end

  result

end

def md5sum( filename )
  result = safe_backticks( $md5sum_binary, filename )
  if result =~ /^([a-z0-9]{32})\s/
    $1
  else
    nil
  end
end

# This class holds details of single difficulty level and steps type
# based on information in the DWI file.  (It should correspond to an
# entry in the Stats.xml file, though.)

class Steps

  attr_accessor :steps_type, :difficulty, :feet
  attr_accessor :cleared, :played
  attr_accessor :song

  def initialize( steps_type, difficulty, feet, song )
    @steps_type = steps_type
    @difficulty = difficulty
    @feet = feet
    @high_scores = Array.new
    @cleared = 0
    @played = 0
    @song = song
  end

  def add_high_score( score )
    @high_scores.push score
    @high_scores = @high_scores.sort.reverse
  end

  def add_high_scores( score_list )
    @high_scores = @high_scores.append( score_list )
    @high_scores = @high_scores.sort.reverse
  end

  def highest_score
    if @high_scores.length > 0
      return @high_scores[0].to_i
    else
      return 0
    end
  end

  def xml_steps_type
    case @steps_type
    when 'SINGLE'
      return 'dance-single'
    when 'DOUBLE'
      return 'dance-double'
    else
      raise "Unknown type of steps: " + @steps_type
    end     
  end

  def xml_difficulty
    case @difficulty
    when 'BEGINNER'
      return 'Beginner'
    when 'BASIC'
      return 'Easy'
    when 'ANOTHER'
      return 'Medium'
    when 'MANIAC'
      return 'Hard'
    when 'SMANIAC'
      return 'Challenge'
    else
      raise "Unknown difficulty level: " + @difficulty
    end
  end

  def to_s
    "(" + feet.to_s + ") [" + xml_difficulty + "] {" + steps_type[0..0] + '} "' + song.leaf_name + '"'
  end

  def to_tr
    "<tr>" +
      "<td class=\"high-score-cell image\"><img src=\"#{song.key}.png\" alt=\"Banner for #{song.leaf_name}\" title=\"Banner for #{song.leaf_name}\"></td>" +
      "<td class=\"high-score-cell score\">#{highest_score.to_s}</td>" +
      "<td class=\"high-score-cell played\">#{played.to_s}</td>" +
      "<td class=\"high-score-cell cleared\">#{cleared.to_s}</td>" +
      "<td class=\"high-score-cell difficulty\">#{xml_difficulty}</td>" +
      "<td class=\"high-score-cell feet\">#{feet.to_s}</td>" +
      "<td class=\"high-score-cell steps-type\">#{steps_type}</td>" +
      "<td class=\"high-score-cell song-name\">#{song.leaf_name}</td>" +
      "</tr>"
  end

  def Steps.tr_header
    "<tr>" +
      "<th class=\"high-score-cell image-header\">Image</th>" +
      "<th class=\"high-score-cell score-header\">Best Score</th>" +
      "<th class=\"high-score-cell played-header\">Times Played</th>" +
      "<th class=\"high-score-cell cleared-header\">Times Cleared</th>" +
      "<th class=\"high-score-cell difficulty-header\">Difficulty</th>" +
      "<th class=\"high-score-cell feet-header\">Feet</th>" +
      "<th class=\"high-score-cell steps-type-header\">Steps Type</th>" +
      "<th class=\"high-score-cell song-name-header\">Song Name</th>" +
      "</tr>"
  end

end

class HighScore
  
  attr_reader :marvellous
  attr_reader :perfect
  attr_reader :great
  attr_reader :good
  attr_reader :boo
  attr_reader :miss

  attr_reader :ok
  attr_reader :ng

  attr_reader :score

  def initialize( high_score_element, location )

    @score = Integer( high_score_element.elements['Score'].text )

    @marvellous = Integer( high_score_element.elements['TapNoteScores/Marvelous'].text )
    @perfect = Integer( high_score_element.elements['TapNoteScores/Perfect'].text )
    @great = Integer( high_score_element.elements['TapNoteScores/Great'].text )
    @good = Integer( high_score_element.elements['TapNoteScores/Good'].text )
    @boo = Integer( high_score_element.elements['TapNoteScores/Boo'].text )
    @miss = Integer( high_score_element.elements['TapNoteScores/Miss'].text )
    
    @ok = Integer( high_score_element.elements['HoldNoteScores/OK'].text )
    @ng = Integer( high_score_element.elements['HoldNoteScores/NG'].text )
    
  end

  def to_i
    @score
  end

  def <=>( other )
    self.to_i <=> other.to_i
  end

end

class SongLocation

  attr_reader :stepmania_root
  attr_reader :directory
  attr_reader :mp3_file_name
  attr_reader :dwi_file_name
  attr_reader :png_file_name
  attr_reader :leaf_name
  attr_reader :key

  def to_s
    @directory
  end

  def exclude_song?

    # This is a particularly stupid way to exclude songs that I don't
    # want to appear in the high score tables...

    if leaf_name =~ /OOPS.*DID IT AGAIN/
      return true
    elsif leaf_name =~ /HOT LIMIT/
      return true
    elsif directory =~ /Disney/i
      return true
    else
      return false
    end

  end

  def initialize( stepmania_root, song_element )

    @stepmania_root = stepmania_root
    @song_element = song_element
    @directory = @song_element.attributes["Dir"]
    
    @leaf_name = nil

    @mp3_file_name = nil
    @dwi_file_name = nil
    @png_file_name = nil

    @song_root = nil
    @key = nil

    if directory =~ Regexp.new( '([^\/]+)\/?$' )

      @leaf_name = $1
      @song_root = $stepmania_root_directory_name + directory + "/"

      @mp3_file_name = @song_root + @leaf_name + ".mp3"
      @dwi_file_name = @song_root + @leaf_name + ".dwi"
      @png_file_name = @song_root + @leaf_name + ".png"

    else

      raise "The directory name " + @directory + "was malformed"

    end

    if FileTest.exist?( @mp3_file_name ) && FileTest.exist?( @dwi_file_name )
      @key = md5sum( @mp3_file_name ) + md5sum( @dwi_file_name )
    else
      @key = nil
    end

  end

  def get_steps_from_dwi
    
    steps_from_dwi = Array.new

    if FileTest.exist?( @dwi_file_name )
      open @dwi_file_name, "r" do |dwi_file|
        dwi_file.each do |line|
          line.chomp!
          if( line =~ /^#(SINGLE|DOUBLE):(BEGINNER|BASIC|ANOTHER|MANIAC|SMANIAC):([0-9]+):/ )
            steps_from_dwi.push Steps.new( $1, $2, Integer($3), self )
          end
        end
      end
    end

    steps_from_dwi

  end

  def get_high_scores_and_add_to( steps_from_dwi )

    @song_element.each do |steps_element|

      if steps_element.class == REXML::Element && steps_element.name =="Steps"
        
        steps_type = steps_element.attributes['StepsType']
        difficulty = steps_element.attributes['Difficulty']

        $verbose && STDERR.puts( "  " + difficulty + " (" + steps_type + ")" )

        high_score_list_element = steps_element.elements['HighScoreList']

        number_of_times_played = Integer( high_score_list_element.elements['NumTimesPlayed'].text )

        dwi_steps_object = nil

        steps_from_dwi.each do |s|
          next if (s.xml_steps_type != steps_type)
          next if (s.xml_difficulty != difficulty)
          dwi_steps_object = s
        end
        
        unless dwi_steps_object
          $verbose && STDERR.puts( "Couldn't find steps of the right type and difficulty in the DWI file" )
          $verbose && STDERR.puts( "(for directory \"" + @leaf_name + "\")" )
          next
        end
        
        number_of_times_cleared = 0

        high_score_list_element.each do |high_score_element|
          if high_score_element.class == REXML::Element && high_score_element.name == "HighScore"
            high_score = HighScore.new( high_score_element, self )

            $verbose && STDERR.puts( "      " + high_score.to_i.to_s )
            if dwi_steps_object
              dwi_steps_object.add_high_score( high_score )
              number_of_times_cleared += 1
            end
            
          end
        end

        dwi_steps_object.played += number_of_times_played
        $verbose && STDERR.puts( "    played " + number_of_times_played.to_s + " times" )
        dwi_steps_object.cleared += number_of_times_cleared
        $verbose && STDERR.puts( "    cleared " + number_of_times_cleared.to_s + " times" )

      end
    end

  end

end

# This class encapsulates a Song; it might correspond to many Song
# elements from various Stats.xml files.

class Song

  attr_reader :key

  attr_reader :steps_from_dwi

  attr_accessor :locations

  def initialize( location )

    @key = location.key
    @steps_from_dwi = location.get_steps_from_dwi
    @locations = [ location ]

    location.get_high_scores_and_add_to( @steps_from_dwi )

  end

  def exclude?
    return locations.detect { |l| l.exclude_song? }
  end

  def add_location( new_location )

    unless key
      raise "[BUG] self.key is nil in add_location"
    end
    unless new_location.key
      raise "[BUG] new_location.key is nil in add_location"
    end
    unless key == new_location.key
      raise "[BUG] keys don't match in new_location"
    end
    
    @locations.push new_location

    new_location.get_high_scores_and_add_to( @steps_from_dwi )

  end

end

songs_hash = Hash.new
keys_of_images = Hash.new

open( output_index_name, "w" ) do |o|

  open( file_name, "r" ) do |f|

    STDERR.print "About to read Stats.xml file..."
    STDERR.flush
  
    document = REXML::Document.new( f )

    STDERR.puts " done."
  
    # That's the bit that takes all the time, now we can do whatever we
    # like...
  
    STDERR.print "Creating Song objects from XML data..."
    STDERR.flush

    document.elements.each("Stats/SongScores/Song") do |e|

      song_location = SongLocation.new( $stepmania_root_directory_name, e )

      $verbose && STDERR.puts( "Considering: " + song_location.to_s )

      key = song_location.key

      if key && ! song_location.exclude_song?

        if songs_hash.has_key?( key )
          songs_hash[key].add_location( song_location )
        else
          songs_hash[key] = Song.new( song_location )
        end
    
        $verbose && STDERR.puts( "  So that's: " + songs_hash[key].locations[0].leaf_name )
        $verbose && STDERR.puts( "  ... in " + songs_hash[key].locations.length.to_s + " locations" )

      else
        $verbose && STDERR.puts( "  Ignoring directory: " + song_location.directory )
        $verbose && STDERR.puts( "  (It's missing the MP3 or DWI file.)" )
      end
  
    end
  
    STDERR.puts " done."

    STDERR.print "Sorting high scores into tables..."
    STDERR.flush

    different_feet_single = Hash.new
    different_feet_double = Hash.new
  
    uncleared_by_feet_single = Hash.new
    uncleared_by_feet_double = Hash.new
  
    all_steps = Array.new
      
    songs_hash.each_value do |song|
      song.steps_from_dwi.each do |s|
          
        if $include_doubles || (s.steps_type == "SINGLE")
          all_steps.push s
        end
  
        if s.cleared == 0
  
          keys_of_images[song.key] = true
  
          if s.steps_type == "SINGLE"
            if uncleared_by_feet_single.has_key? s.feet
              uncleared_by_feet_single[s.feet].push( s )
            else
              uncleared_by_feet_single[s.feet] = [ s ]
            end
          elsif s.steps_type == "DOUBLE" && $include_doubles
            if uncleared_by_feet_double.has_key? s.feet
              uncleared_by_feet_double[s.feet].push( s )
            else
              uncleared_by_feet_double[s.feet] = [ s ]
            end
          end
          
        end
        
        if s.highest_score > 0
          
          keys_of_images[song.key] = true
          
          if s.steps_type == "SINGLE"
            if different_feet_single.has_key? s.feet
              different_feet_single[s.feet].push( s )
            else
              different_feet_single[s.feet] = [ s ]
            end
          elsif s.steps_type == "DOUBLE" && $include_doubles
            if different_feet_double.has_key? s.feet
              different_feet_double[s.feet].push( s )
            else
              different_feet_double[s.feet] = [ s ]
            end
          end
          
          end
      end
    end
  
    different_feet_single_sorted = different_feet_single.keys.sort.reverse
    different_feet_double_sorted = different_feet_double.keys.sort.reverse
  
    uncleared_by_feet_single_sorted = uncleared_by_feet_single.keys.sort.reverse.find_all { |i|
        (i >= $minimum_steps) && (i <= $maximum_steps) }
    uncleared_by_feet_double_sorted = uncleared_by_feet_double.keys.sort.reverse.find_all { |i|
        (i >= $minimum_steps) && (i <= $maximum_steps) }
  
    all_steps = all_steps.sort { |a,b| a.highest_score <=> b.highest_score }.reverse
  
    most_played = all_steps.sort { |a,b| a.played <=> b.played }.reverse

    STDERR.puts " done."
  
    STDERR.print "Creating HTML index file..."
    STDERR.flush

    $verbose && STDERR.puts( "Best scores overall:" )
    o.puts "<div class=\"best-scores\">"
    o.puts "<h3 class=\"high-score-table-title\"><a name=\"best-overall\">Best Scores Overall</a></h3>"
    o.puts "<table class=\"high-score-table\">"  
    o.puts Steps.tr_header
  
    all_steps[0...20].each do |s|
      keys_of_images[s.song.key] = true          
      $verbose && STDERR.puts( sprintf( "%12d ", s.highest_score ) + s.to_s )
      o.puts s.to_tr
    end
  
    o.puts "</table></div>"
  
    different_feet_single_sorted.each do |feet|
  
      o.puts "<div class=\"best-scores\">"
      o.puts "<h3 class=\"high-score-table-title\"><a name=\"best-#{feet.to_s}-feet-single\">Best #{feet} Feet Songs (SINGLE):</a></h3>"
      o.puts "<table class=\"high-score-table\">"  
      o.puts Steps.tr_header
      $verbose && STDERR.puts( "Best #{feet} Feet Songs (SINGLE):" )
  
      feet_specific_steps = different_feet_single[feet].sort { |a,b|
          a.highest_score <=> b.highest_score }.reverse
  
      feet_specific_steps[0...10].each do |s|
        keys_of_images[s.song.key] = true          
        $verbose && STDERR.puts( sprintf( "%12d ", s.highest_score ) + s.to_s )
        o.puts s.to_tr
      end
  
      o.puts "</table></div>"
  
    end

    if $include_doubles  

      different_feet_double_sorted.each do |feet|
    
        o.puts "<div class=\"best-scores\">"
        o.puts "<h3 class=\"high-score-table-title\"><a name=\"best-#{feet.to_s}-feet-double\">Best #{feet} Feet Songs (DOUBLE):</a></h3>"
        o.puts "<table class=\"high-score-table\">"
        o.puts Steps.tr_header
        $verbose && STDERR.puts( "Best #{feet} Feet Songs (DOUBLE):" )
    
        feet_specific_steps = different_feet_double[feet].sort { |a,b|
            a.highest_score <=> b.highest_score }.reverse
    
        feet_specific_steps[0...10].each do |s|
          keys_of_images[s.song.key] = true          
          $verbose && STDERR.puts( sprintf( "%12d ", s.highest_score ) + s.to_s )
          o.puts s.to_tr
        end
    
        o.puts "</table></div>"
    
      end

    end
  
    o.puts "<div class=\"best-scores\">"
    o.puts "<h3 class=\"high-score-table-title\"><a name=\"most-played\">Most Played Songs</a></h3>"
    o.puts "<table class=\"high-score-table\">"
    o.puts Steps.tr_header
    $verbose && STDERR.puts( "Most played songs: " )
  
    most_played[0...20].each do |s|
      keys_of_images[s.song.key] = true          
      $verbose && STDERR.puts( sprintf( "%12d ", s.played ) + s.to_s )
      o.puts s.to_tr
    end
  
    o.puts "</table></div>"
  
    uncleared_by_feet_single_sorted.each do |feet|
  
      o.puts "<div class=\"best-scores\">"
      o.puts "<h3 class=\"high-score-table-title\"><a name=\"uncleared-#{feet.to_s}-feet-single\">Uncleared #{feet} Feet Songs (SINGLE)</a></h3>"
      o.puts "<table class=\"high-score-table\">"
      o.puts Steps.tr_header
  
      $verbose && STDERR.puts( "Uncleared #{feet} Feet songs (SINGLE):" )
  
      feet_specific_steps = uncleared_by_feet_single[feet].sort { |a,b|
          a.played <=> b.played }.reverse
  
      feet_specific_steps[0...20].each do |s|
        keys_of_images[s.song.key] = true          
        $verbose && STDERR.puts( sprintf( "%12d ", s.played ) + s.to_s )
        o.puts s.to_tr
      end
   
      o.puts "</table></div>"
  
    end
  
    if $include_doubles

      uncleared_by_feet_double_sorted.each do |feet|
    
        o.puts "<div class=\"best-scores\">"
        o.puts "<h3 class=\"high-score-table-title\"><a name=\"uncleared-#{feet.to_s}-feet-double\">Uncleared #{feet} Feet Songs (DOUBLE)</a></h3>"
        o.puts "<table class=\"high-score-table\">"
        o.puts Steps.tr_header
    
        $verbose && STDERR.puts( "Uncleared #{feet} Feet songs (DOUBLE):" )
    
        feet_specific_steps = uncleared_by_feet_double[feet].sort { |a,b|
            a.played <=> b.played }.reverse
    
        feet_specific_steps[0...20].each do |s|
          keys_of_images[s.song.key] = true          
          $verbose && STDERR.puts( sprintf( "%12d ", s.played ) + s.to_s )
          o.puts s.to_tr
        end
    
        o.puts "</table></div>"
  
      end

    end
  
    STDERR.puts " done."

  end

end

STDERR.print "Creating thumbnail images..."
STDERR.flush
  
keys_of_images.each_key do |key|

  song = songs_hash[key]
  
  if FileTest.exist? song.locations[0].png_file_name

    system( $convert_binary,
            "-scale",
            "128x40",
            song.locations[0].png_file_name,
            output_directory_name + "/" + key + ".png" )
    
  end

end

STDERR.puts " done."