#!/usr/bin/env ruby

# RailRoad - RoR diagrams generator
# http://railroad.rubyforge.org/
#
# Copyright 2007 - Javier Smaldone (http://www.smaldone.com.ar)
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# Purpose:
#   RailRoad generates models and controllers diagrams in DOT language
#   for a Rails application.

@app_name  = "railroad"
@app_human_name = "RailRoad"
@version   = "0.1.2"
@copyright = "Copyright (C) 2007 Javier Smaldone"

def print_help
  print "Usage: #{@app_name} [options] [command]\n\n"
  print "Options:\n"
  print "  -a, --all\tInclude all models (not ActiveRecord::Base derived)\n"
  print "  -b, --brief\tGenerate compact diagram (no attributes)\n"
  print "  -c, --compact\tConcentrate edges (only for modules diagram)\n"
  print "  -l, --label\tAdd a label with diagram info (type, date, version)\n"
  print "  -i, --inherit\tInclude inheritance relations\n"
  print "  -h, --help\tDisplay this help message and exit\n"
  print "  -v, --version\tDisplay version information and exit\n\n"
  print "Commands:\n"
  print "  -M\tGenerate models diagram\n"
  print "  -C\tGenerate controllers diagram\n"
  print "\n"
end

def node(name)
    '"' + name + '"'
end

def print_info(type)
  print "\t_diagram_info [shape=\"plaintext\", label=\"Diagram: #{type}\\l"
  print "Date: #{Time.now.strftime "%b %d %Y - %H:%M"}\\l"
  print "Migration version: #{ActiveRecord::Migrator.current_version}\\l"
  print "Generated by #{@app_human_name} #{@version}\\l\""
  print ", fontsize=14]\n"
end

def print_error(type)
    STDERR.print "Error loading #{type}. (Are you running "
    STDERR.print "#{@app_name} on the aplication's root directory?)\n\n"
end

def generate_models
  # Load model classes
  begin
    Dir.glob("app/models/**/*.rb") {|m| require m }
  rescue LoadError
    print_error "model classes"
    raise
  end

  # DOT diagram header
  print "digraph models_diagram {\n"
  print "\tgraph[concentrate=true]\n" if @compact
  print_info "Models" if @label

  models_files = Dir.glob("app/models/**/*.rb")
  models_files.each do |m|
    # Extract the class name from the file name
    current_class = m.split('/')[2..-1].join('/').split('.').first.camelize.constantize

    # Is current_clas derived from ActiveRecord::Base?
    if current_class.respond_to?'reflect_on_all_associations'

      # Print the node
      if @brief
        print "\t#{node(current_class.name)}\n"
      else 
        print "\t#{node(current_class.name)} "
        print "[shape=Mrecord, label=\"{#{current_class.name}|"
        current_class.content_columns.each do |a|
          print "#{a.name} :#{a.type.to_s}\\l" 
        end
        print "}\"]\n"
      end

      # Iterate over the class associations
      current_class.reflect_on_all_associations.each do |association|

        # Only non standard association names needs a label
        if association.class_name != association.name.to_s.singularize.camelize
          params = "label=\"#{association.name.to_s}\", "
        else
          params = ""
        end

        # Skip "belongs_to" associations
        if association.macro.to_s != 'belongs_to'
          # Put a tail label with the association arity
          if association.macro.to_s == 'has_one' 
            params += 'taillabel="1"'
          else
            params += 'taillabel="n"'
          end

          # Print the edge
          print "\t\t#{node(current_class.name)} -> "
          print "#{node(association.class_name)} "
          print "[ #{params} ]\n"
        end
      end # Associations

      if @inherit && (current_class.superclass != ActiveRecord::Base)
        # Print the inheritance edge
        print "\t\t#{node(current_class.name)} -> "
	print "#{node(current_class.superclass.name)}"
        print " [arrowhead=\"onormal\"]\n"
      end      

    elsif @all_classes
      # Not ActiveRecord::Base model
      if @brief
        print "\t#{node(current_class.name)} [shape=box]\n"
      else 
        print "\t#{node(current_class.name)}"
        print "[shape=record, label=\"{#{current_class.name}|}\"]\n"
      end
    end
  end # Dir
  print "}\n"
end # generate_models

def generate_controllers
  # Load controller classes
  begin
    # ApplicationController must be loaded first
    require "app/controllers/application.rb" 
    Dir.glob("app/controllers/**/*_controller.rb") do |c| 
      require c
    end
  rescue LoadError
    print_error "controller classes"
    raise
  end

  # DOT diagram header
  print "digraph controllers_diagram {\n"
  print "\tgraph[overlap=false, splines=true]\n"
  print_info "Controllers" if @label

  app_controller = ApplicationController

  controllers_files = Dir.glob("app/controllers/**/*_controller.rb")
  controllers_files << 'app/controllers/application.rb'
  controllers_files.each do |c|
    # Extract the class name from the file name
    class_name = c.split('/')[2..-1].join('/').split('.').first.camelize

    # ApplicationController's file is 'application.rb'
    if class_name == 'Application' 
      class_name = 'ApplicationController' 
    end 
    current_class = class_name.constantize

    # Print the node
    if @brief
      print "\t#{node(current_class.name)}\n"
    else 
      print "\t#{node(current_class.name)} "
      print "[shape=Mrecord, label=\"{#{current_class.name}|"
      current_class.public_instance_methods(false).sort.each do |m|
        print "#{m}\\l" 
      end
      print "|"
      current_class.protected_instance_methods(false).sort.each do |m|
        print "#{m}\\l" 
      end
      print "|"
      current_class.private_instance_methods(false).sort.each do |m|
        print "#{m}\\l" 
      end
      print "}\"]\n"
    end
    if @inherit && (app_controller.subclasses.include? current_class.name)
      # Print the inheritance edge
      print "\t#{node(current_class.name)} -> "
      print "#{node(current_class.superclass.name)} "
      print "[arrowhead=\"onormal\"]\n"
    end
  end # Dir
  print "}\n"
end # generate_controllers

# Main program 
if ARGV.include?('--help') || ARGV.include?('-h')
  print_help
  exit
elsif ARGV.include?('--version') || ARGV.include?('-v')
  print "#{@app_name} #{@version}\n#{@copyright}\n\n"
  exit
end

@all_classes = ARGV.include?('--all') || ARGV.include?('-a')
@brief = ARGV.include?('--brief') || ARGV.include?('-b')
@compact = ARGV.include?('--compact') || ARGV.include?('-c')
@label = ARGV.include?('--label') || ARGV.include?('-l')
@inherit = ARGV.include?('--inherit') || ARGV.include?('-i')

params = %w(-a --all -b --brief -c --compact -l --label -i --inherit -M -C)
error_params = ARGV - params
if ! error_params.empty? 
  # Unknown parameter
  STDERR.print "Error: wrong parameter(s) #{error_params}\n"
  STDERR.print "(try '#{@app_name} --help')\n\n"
  exit 1
end
if ARGV.include?('-M') && ARGV.include?('-C')
  STDERR.print "Error: you can only specify one command\n"
  STDERR.print "(try '#{@app_name} --help')\n\n"
  exit 1
elsif ! (ARGV.include?('-M') || ARGV.include?('-C'))
  STDERR.print "Error: you must specify a command\n"
  STDERR.print "(try '#{@app_name} --help')\n\n"
  exit 1
end

# Load Rails app environment
begin
  require "config/environment"
rescue LoadError
  print_error "app environment"
  raise
end

if ARGV.include?('-M')
  generate_models
elsif ARGV.include?('-C')
  generate_controllers
end