#!/usr/bin/env ruby

# RailRoad - RoR diagrams generator
# http://railroad.rubyforge.org
#
#   RailRoad generates models and controllers diagrams in DOT language
#   for a Rails application.
#
# 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.
#

APP_NAME  = "railroad"
APP_HUMAN_NAME = "RailRoad"
APP_VERSION   = [0,3,3]
COPYRIGHT = "Copyright (C) 2007 Javier Smaldone"

require 'ostruct'

# Command line options structure
class OptionsStruct < OpenStruct

  require 'optparse'

  def initialize
    init_options = { :all => false,
                     :brief => false,
                     :inheritance => false,
                     :join => false,
                     :label => false,
                     :modules => false,
                     :hide_types => false,
                     :hide_public => false,
                     :hide_protected => false,
                     :hide_private => false,
                     :command => '' }
    super(init_options)
  end # initialize

  def parse(args)
    @opt_parser = OptionParser.new do |opts|
      opts.banner = "Usage: #{APP_NAME} [options] command"
      opts.separator ""
      opts.separator "Common options:"
      opts.on("-b", "--brief", "Generate compact diagram", 
              "  (no attributes nor methods)") do |b|
        self.brief = b
      end
      opts.on("-i", "--inheritance", "Include inheritance relations") do |i|
        self.inheritance = i
      end
      opts.on("-l", "--label", "Add a label with diagram information", 
              "  (type, date, migration, version)") do |l|
        self.label = l
      end
      opts.on("-o", "--output FILE", "Write diagram to file FILE") do |f|
        self.output = f
      end
      opts.separator ""
      opts.separator "Models diagram options:"
      opts.on("-a", "--all", "Include all models", 
              "  (not only ActiveRecord::Base derived)") do |a|
        self.all = a
      end
      opts.on("--hide-types", "Hide attributes type") do |h|
        self.hide_types = h
      end
      opts.on("-j", "--join", "Concentrate edges") do |j|
        self.join = j
      end
      opts.on("-m", "--modules", "Include modules") do |m|
        self.modules = m
      end
      opts.separator ""
      opts.separator "Controllers diagram options:"
      opts.on("--hide-public", "Hide public methods") do |h|
        self.hide_public = h
      end
      opts.on("--hide-protected", "Hide protected methods") do |h|
        @options.hide_protected = h
      end
      opts.on("--hide-private", "Hide private methods") do |h|
        @options.hide_private = h
      end
      opts.separator ""
      opts.separator "Other options:"
      opts.on("-h", "--help", "Show this message") do
        print "#{opts}\n"
        exit
      end
      opts.on("--version", "Show version and copyright") do
        print "#{APP_HUMAN_NAME} version #{APP_VERSION.join('.')}\n\n" +
              "#{COPYRIGHT}\n" +
              "This is free software; see the source for copying conditions.\n\n"
        exit
      end
      opts.separator ""
      opts.separator "Commands (you must supply one of these):"
      opts.on("-M", "--models", "Generate models diagram") do |c|
         if self.command == 'controllers'
           STDERR.print "Error: Can't generate models AND " + 
                        "controllers diagram\n\n"
           exit 1
         else 
           self.command = 'models'        
         end
      end
      opts.on("-C", "--controllers", "Generate controllers diagram") do |c|
         if self.command == 'models'
           STDERR.print "Error: Can't generate models AND " +
                        "controllers diagram\n\n"
           exit 1
         else 
           self.command = 'controllers'        
         end
      end
      opts.separator ""
      opts.separator "For bug reporting and additional information, please see:"
      opts.separator "http://railroad.rubyforge.org"
    end

    begin
      @opt_parser.parse!(args)
    rescue OptionParser::AmbiguousOption
      option_error "Ambiguous option"
    rescue OptionParser::InvalidOption
      option_error "Invalid option"
    rescue OptionParser::InvalidArgument
      option_error "Invalid argument"
    rescue OptionParser::MissingArgument
      option_error "Missing argument"
    rescue
      option_error "Unknown error"
    end
  end  # initialize

  private 

  def option_error(msg)
    STDERR.print "Error: #{msg}\n\n #{@opt_parser}\n"
    exit 1
  end

end  # class OptionsStruct


# Root class for RailRoad diagrams
class AppDiagram

  def initialize(options)
    @options = options
    load_environment
    load_classes
  end

  private 

  # Quotes a class name
  def node(name)
    '"' + name + '"'
  end

  # Prevents Rails application from writing to STDOUT
  def disable_stdout
    @old_stdout = STDOUT.dup
    STDOUT.reopen( PLATFORM =~ /mswin/ ? "NUL" : "/dev/null" )
  end

  # Restore STDOUT  
  def enable_stdout
    STDOUT.reopen(@old_stdout)
  end

  # Print diagram header
  def print_header(type)
    print "digraph #{type.downcase}_diagram {\n" +
          "\tgraph[overlap=false, splines=true]\n"
    print_info(type) if @options.label
  end
  

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

  # Print an edge
  def print_edge(from, to, attributes)
    print "\t#{node(from)} -> #{node(to)} [#{attributes}]\n"
  end

  # Print a node
  def print_node(name, attributes=nil)
    print "\t#{node(name)}"
    print " [#{attributes}]" if attributes && attributes != ''
    print "\n"
  end

  # Print error when loading Rails application
  def print_error(type)
    STDERR.print "Error loading #{type}.\n  (Are you running " +
                 "#{APP_NAME} on the aplication's root directory?)\n\n"
  end

  # Load Rails application's environment
  def load_environment
    begin
      disable_stdout
      require "config/environment"
      enable_stdout
    rescue LoadError
      enable_stdout
      print_error "application environment"
      raise
    end
  end

end # class AppDiagram


# RailRoad models diagram
class ModelsDiagram < AppDiagram

  def initialize(options)
    super options 
    # Processed habtm associations
    @habtm = []
  end

  # Generate models diagram
  def generate
    print_header 'Models'
    models_files = Dir.glob("app/models/**/*.rb")
    models_files.each do |m|
      # Extract the class name from the file name
      class_name = m.split('/')[2..-1].join('/').split('.').first.camelize
      process_class class_name.constantize  
    end 
    print "}\n"
  end # generate

  private

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

  # Process model class
  def process_class(current_class)

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

      # Print the node
      if @options.brief
        node_attrib = ''
      else 
        node_attrib = 'shape=Mrecord, label="{' + current_class.name + '|'
        current_class.content_columns.each do |a|
          node_attrib += a.name
          node_attrib += ' :' + a.type.to_s unless @options.hide_types
          node_attrib += '\l'
        end
        node_attrib += '}"'
      end
      print_node current_class.name, node_attrib
      class_printed = true
      # Iterate over the class associations
      current_class.reflect_on_all_associations.each do |a|
        process_association(current_class, a)
      end
    elsif @options.all && (current_class.is_a? Class)
      # Not ActiveRecord::Base model
      if @options.brief
        node_attrib = 'shape=box'
      else 
        node_attrib = 'shape=record, label="{' + current_class.name + '|}"'
      end
      print_node current_class.name, node_attrib
      class_printed = true
    elsif @options.modules && (current_class.is_a? Module)
        print_node current_class.name, 
                   'shape=box, style=dotted, label="' + current_class.name + '"'        
    end

    # Only consider meaningful inheritance relations for printed classes
    if @options.inheritance && class_printed && 
       (current_class.superclass != ActiveRecord::Base) &&
       (current_class.superclass != Object)
      print_edge(current_class.name, current_class.superclass.name, 
                 'arrowhead="onormal"')
    end      

  end # process_class

  # Process model association
  def process_association(current_class, association)

    # Skip "belongs_to" associations
    return if association.macro.to_s == 'belongs_to'
 
    assoc_attrib = ""
    assoc_name = ""
    # Only non standard association names needs a label
    if association.class_name != association.name.to_s.singularize.camelize
      association_name = association.name.to_s
      assoc_attrib += 'label="' + association_name + '", '
    end

    # Tail label with the association arity
    assoc_attrib += association.macro.to_s == 'has_one' ? 'taillabel="1"' : 'taillabel="n"'

    # Use double-headed arrows for habtm and has_many, :through
    if association.macro.to_s == 'has_and_belongs_to_many' ||
       (association.macro.to_s == 'has_many' && association.options[:through])
      
      # Skip habtm associations already considered
      return if @habtm.include? [association.class_name, current_class.name, 
                          association_name]
      @habtm << [current_class.name, association.class_name,association_name]
      assoc_attrib += ', headlabel="n", arrowtail="normal"'
    end  
    print_edge(current_class.name, association.class_name, assoc_attrib)
  end # process_association

end # class ModelsDiagram


# RailRoad controllers diagram
class ControllersDiagram < AppDiagram
 
  def initialize(options)
    super options
    @app_controller = ApplicationController
  end

  # Generate controllers diagram
  def generate

    print_header 'Controllers'

    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'
      class_name += 'Controller' if class_name == 'Application'
      process_class class_name.constantize
    end 
    print "}\n"
  end # generate

  private

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

  # Proccess controller class
  def process_class(current_class)

    if @options.brief
      print_node  current_class.name

    elsif current_class.is_a? Class 
      node_attrib = 'shape=Mrecord, label="{' + current_class.name + '|'
      current_class.public_instance_methods(false).sort.each { |m|
        node_attrib += m + '\l' 
      } unless @options.hide_public
      node_attrib += '|'
      current_class.protected_instance_methods(false).sort.each { |m|
        node_attrib += m + '\l'
      } unless @options.hide_protected
      node_attrib += '|'
      current_class.private_instance_methods(false).sort.each { |m|
        node_attrib += m + '\l'
      } unless @options.hide_private
      node_attrib += '}"' 
      print_node current_class.name, node_attrib 

    elsif @options.modules && current_class.is_a?(Module)
        print_node current_class.name, 
                   'shape=box, style=dotted, label="' + current_class.name + '"'
    end

    # Print the inheritance edge (only for ApplicationControllers)
    if @options.inheritance && 
       (@app_controller.subclasses.include? current_class.name)
      print_edge(current_class.name, current_class.superclass.name, 
                 'arrowhead="onormal"')
    end

  end # process_class

end # class ControllersDiagram


# Main program

options = OptionsStruct.new

options.parse ARGV

if options.command == 'models'
  diagram = ModelsDiagram.new options
elsif options.command == 'controllers'
  diagram = ControllersDiagram.new options
else
  STDERR.print "Error: You must supply a command\n" +
               "  (try #{APP_NAME} -h)\n\n"
  exit 1
end

if options.output
  old_stdout = STDOUT.dup
  begin
     STDOUT.reopen(options.output)
  rescue
     STDERR.print "Error: Cannot write diagram to #{options.output}\n\n"
     exit 2
  end
end

diagram.generate

if options.output
  STDOUT.reopen(old_stdout)
end