#!/usr/bin/env ruby
#---
# Excerpted from "Agile Web Development with Rails"
# We make no guarantees that this code is fit for any purpose. 
# Visit http://www.pragmaticprogrammer.com for more book information.
#---
#
# Copyright (C) 2005 Michael Schubert
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
#  "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.   

# Usage: rails_graph.rb <path to recursively search for .rb files>

class GraphClass
  attr_accessor :name, :parent, :attributes, :methods
    
  def initialize(name)
    @name = name
  end
  
  def split(symbol)
    o = String.new
    symbol.each { |elements| o << elements.to_s.concat("\\n") }
    o      
  end

  def output
    "  \"#{@name}\" [label=\"#{@name}|#{split(@attributes)}|#{split(@methods)}\"];\n"
  end  
end

class GraphPackage
  attr_accessor :name, :shape, :style, :color
  
  def initialize(name)
    @name = name
    @shape = "box"
    @color = "palegoldenrod"
    @style = "filled"    
  end
  
  def output
    attributes = "label=\"#{@name}\", shape=\"#{@shape}\", style=\"#{@style}\", color=\"#{@color}\"" 
    "  \"#{@name}\" [#{attributes}];\n"
  end
  
end

class GraphRelation
  attr_accessor :child, :parent, :type, :p_label, :c_label
  
  def initialize(child, parent)
    @child = child
    @parent = parent
  end
  
  def output
    determine_edge
    attributes = "shape=\"#{@shape}\", arrowhead=\"#{@arrowhead}\""
    attributes << " taillabel=\"#{p_label}\"" if @p_label
    attributes << " headlabel=\"#{c_label}\"" if @c_label
    "  \"#{@child}\" #{@edge} \"#{@parent}\" [#{attributes}]\n"
  end
  
  def determine_edge
    case @type
      when "inheritance"
        @edge = "->"
        @arrowhead = "empty"
        @shape = "normal"
      when "multiplicity"
        @edge = "->"
        @arrowhead = "open"
        @shape = "vee"
        @taillabel = @p_label
        @headlabel = @c_label
      else
        @type == nil
        @edge = "-"
        @arrowhead = "empty"
        @shape = "none"
    end
  end
  
end

class RailsGraph
  
  # graph_type = class, sequence or ER
  attr_accessor :graph_name, :graph_type, :output
  
  def initialize
    @output = Array.new
  end
  
  def run(args = [])
    @output << def_graph
    source_files = `find #{args.empty? ? '.' : args.join(' ')} -name "*.rb"`.split("\n")
    source_files.each { |path| generate_graph(path.to_s) }
    @output << end_graph
    output_file = File.new("output.dot", File::CREAT|File::TRUNC|File::WRONLY, 0644) 
    output_file << @output.uniq
  end

  def def_graph
    "digraph #{@graph_name} {\n\n  node[shape=record]\n"
  end

  def end_graph
    "}\n"
  end

  # Where the magic happens
  def node_split(string)
    string = crush_parans(string.sub(/class/,'').sub(/module/,''))
    split_string = string.split('<', 2)
    @node = split_string[0].strip
    if split_string[1]
      @parent = split_string[1].strip
    else
      @parent = nil
    end
  end
  
  # Crush anything between two ()'s to just '()'
  def crush_parans(string)
    # case 1: Foo.method() case 2: Foo.method=() case 3: method()
    "#{string.gsub(/\..*\(.*\)/,'()').gsub(/=.*\(.*\)/,'()').gsub(/\(.*\)/,'()')}"
  end
    
  def method_split(string)
    "#{crush_parans(scrub_record(string.sub(/def\ /,'')))}"
  end
    
  def scrub_record(string)
    # Escape quotes, trash anything after a comment, semicolon
    "#{string.gsub(/\"/,'').sub(/#.*/,'').sub(/;.*/,'')}"
  end
  
  # XXX Holy %#$!# 
  def param_split(string) 
    split_string = string.sub(/#.*/,'').sub(/\(.*/,'').gsub(/@.*&.*/,'').gsub(/\"/,'').sub(/\).*/,'').gsub(/<.*/,'').gsub(/>.*/,'').gsub(/\{.*/,'').gsub(/\..*/,'').gsub(/=.*/,'').gsub(/\|/,'').gsub(/,/,'').strip.sub(/\ /, '\n').gsub(/"/, '\"').gsub(/\ @/,'@').gsub(/\|/, '\|').strip
      "#{split_string}"
  end

  def find_relation_to(string)
    string.scan(/:.*/) { |s| return s.sub(/:/,'').strip.capitalize }
  end
  # XXX End of really bad kludges

  # XXX Break into smaller functions for class, modules, and relationships and output
  def generate_graph(filename)
    methods = Array.new
    parameters = Array.new
    @node = String.new
    @parent = String.new
    non_public = false
    File::open(filename) { |file|
      file.each_line { |line|
        line.scan(/.*class\ .*/) { |class_def|
          node_split(class_def)
        }
        # Scan for Helpers
        line.scan(/module\ .*Helper.*/) { |module_def|
          node_split(module_def)
          @parent = @node.sub(/Helper/,'Controller').strip  # We cheat    
        }
        # Scan for model relationships
        if @parent == "ActiveRecord::Base"
          @plural = false
          line.scan(/belongs_to.*:.*/) { |model_relationship|
            case model_relationship
              when /has_and_belongs_to_many/
                @type = "habtm"
                @plural = true
              when /belongs_to_many/
                @type = "btm"
                @plural = true
              else
                @type = "bt"
            end
            @relation_to = find_relation_to(model_relationship)
          }          
          line.scan(/has_.*:.*/) { |model_relationship|
            case model_relationship
              when /has_many/
                @type = "hm"
                @plural = true
              when /has_one/
                @type = "ho"
            end
            @relation_to = find_relation_to(model_relationship)
          } unless @relation_to
          # For plural ies -> y really ugly
          if @relation_to
            if @relation_to.index('ies') and @plural
                @relation_to = @relation_to.slice(0, @relation_to.index('ies')) + "y"
            else
                @relation_to = @relation_to.chop if @plural
            end
            if @node.index('ies') and @plural
                @node = @node.slice(0, @node.index('ies')) + "y"
            end
            relation = GraphRelation.new(@node, @relation_to)
            relation.type = "multiplicity"
            # Set labels
            case @type
              when "habtm"
                relation.p_label = "0..N"
                relation.c_label = "0..N"
              when "hm"
                relation.p_label = "1"
                relation.c_label = "0..N"
              when "ho"
                relation.c_label = "1"
              when "btm"
                relation.p_label = "0..N"
              when "bt"
                relation.p_label = "1"
            end
            @output << relation.output       
          end
        end
      
        line.scan(/.*private.*/) { |n| non_public = true if n.strip == "private" }
        line.scan(/.*protected.*/) { |n| non_public = true if n.strip == "protected" }
        # If we haven't reached protected or private code, scan for methods and params
        unless non_public
          line.scan(/def\ .*/) { |method_defintion|
            methods << [ method_split(method_defintion) ] 
          }
          methods = methods.uniq  
          line.scan(/@.*/) { |parameter_definition|
            param = param_split(parameter_definition)
            parameters << [ param] unless param == "" or param == "@"
          }
          parameters = parameters.uniq  
        end
      } 
    }
    if @node
      new_node = GraphClass.new(@node)
      new_node.attributes = parameters
      new_node.methods = methods
      @output << new_node.output
      if @parent
        relation = GraphRelation.new(@node, @parent)
        relation.type = "inheritance"
        @output << relation.output
        if @parent.include? "::Base" or @parent.include? "::Observer" or @parent.include? "Struct"
          package = GraphPackage.new(@parent)
          @output << package.output
        end
      end
    end
    
  end
  
end

r = RailsGraph.new
r.graph_name = "RailsApplication"
r.run(ARGV)