#  Created by Luke A. Kanies on 2007-11-07.
#  Copyright (c) 2007. All rights reserved.

require 'puppet/external/dot'
require 'puppet/relationship'
require 'set'

# A hopefully-faster graph class to replace the use of GRATR.
class Puppet::SimpleGraph
  # An internal class for handling a vertex's edges.
  class VertexWrapper
    attr_accessor :in, :out, :vertex

    # Remove all references to everything.
    def clear
      @adjacencies[:in].clear
      @adjacencies[:out].clear
      @vertex = nil
    end

    def initialize(vertex)
      @vertex = vertex
      @adjacencies = {:in => {}, :out => {}}
    end

    # Find adjacent vertices or edges.
    def adjacent(options)
      direction = options[:direction] || :out
      options[:type] ||= :vertices

      return send(direction.to_s + "_edges") if options[:type] == :edges

      @adjacencies[direction].keys.reject { |vertex| @adjacencies[direction][vertex].empty? }
    end

    # Add an edge to our list.
    def add_edge(direction, edge)
      opposite_adjacencies(direction, edge) << edge
    end

    # Return all known edges.
    def edges
      in_edges + out_edges
    end

    # Test whether we share an edge with a given vertex.
    def has_edge?(direction, vertex)
      return(vertex_adjacencies(direction, vertex).length > 0 ? true : false)
    end

    # Create methods for returning the degree and edges.
    [:in, :out].each do |direction|
      # LAK:NOTE If you decide to create methods for directly
      # testing the degree, you'll have to get the values and flatten
      # the results -- you might have duplicate edges, which can give
      # a false impression of what the degree is.  That's just
      # as expensive as just getting the edge list, so I've decided
      # to only add this method.
      define_method("#{direction}_edges") do
        @adjacencies[direction].values.inject([]) { |total, adjacent| total += adjacent.to_a; total }
      end
    end

    # The other vertex in the edge.
    def other_vertex(direction, edge)
      case direction
      when :in; edge.source
      else
        edge.target
      end
    end

    # Remove an edge from our list.  Assumes that we've already checked
    # that the edge is valid.
    def remove_edge(direction, edge)
      opposite_adjacencies(direction, edge).delete(edge)
    end

    def to_s
      vertex.to_s
    end

    private

    # These methods exist so we don't need a Hash with a default proc.

    # Look up the adjacencies for a vertex at the other end of an
    # edge.
    def opposite_adjacencies(direction, edge)
      opposite_vertex = other_vertex(direction, edge)
      vertex_adjacencies(direction, opposite_vertex)
    end

    # Look up the adjacencies for a given vertex.
    def vertex_adjacencies(direction, vertex)
      @adjacencies[direction][vertex] ||= Set.new
      @adjacencies[direction][vertex]
    end
  end

  def initialize
    @vertices = {}
    @edges = []
  end

  # Clear our graph.
  def clear
    @vertices.each { |vertex, wrapper| wrapper.clear }
    @vertices.clear
    @edges.clear
  end

  # Which resources a given resource depends upon.
  def dependents(resource)
    tree_from_vertex(resource).keys
  end

  # Which resources depend upon the given resource.
  def dependencies(resource)
    # Cache the reversal graph, because it's somewhat expensive
    # to create.
    @reversal ||= reversal
    # Strangely, it's significantly faster to search a reversed
    # tree in the :out direction than to search a normal tree
    # in the :in direction.
    @reversal.tree_from_vertex(resource, :out).keys
  end

  # Whether our graph is directed.  Always true.  Used to produce dot files.
  def directed?
    true
  end

  # Determine all of the leaf nodes below a given vertex.
  def leaves(vertex, direction = :out)
    tree = tree_from_vertex(vertex, direction)
    l = tree.keys.find_all { |c| adjacent(c, :direction => direction).empty? }
  end

  # Collect all of the edges that the passed events match.  Returns
  # an array of edges.
  def matching_edges(event, base = nil)
    source = base || event.resource

    unless vertex?(source)
      Puppet.warning "Got an event from invalid vertex #{source.ref}"
      return []
    end
    # Get all of the edges that this vertex should forward events
    # to, which is the same thing as saying all edges directly below
    # This vertex in the graph.
    adjacent(source, :direction => :out, :type => :edges).find_all do |edge|
      edge.match?(event.name)
    end
  end

  # Return a reversed version of this graph.
  def reversal
    result = self.class.new
    vertices.each { |vertex| result.add_vertex(vertex) }
    edges.each do |edge|
      newedge = edge.class.new(edge.target, edge.source, edge.label)
      result.add_edge(newedge)
    end
    result
  end

  # Return the size of the graph.
  def size
    @vertices.length
  end

  # Return the graph as an array.
  def to_a
    @vertices.keys
  end

  # Provide a topological sort.
  def topsort
    degree = {}
    zeros = []
    result = []

    # Collect each of our vertices, with the number of in-edges each has.
    @vertices.each do |name, wrapper|
      edges = wrapper.in_edges
      zeros << wrapper if edges.length == 0
      degree[wrapper.vertex] = edges
    end

    # Iterate over each 0-degree vertex, decrementing the degree of
    # each of its out-edges.
    while wrapper = zeros.pop
      result << wrapper.vertex
      wrapper.out_edges.each do |edge|
        degree[edge.target].delete(edge)
        zeros << @vertices[edge.target] if degree[edge.target].length == 0
      end
    end

    # If we have any vertices left with non-zero in-degrees, then we've found a cycle.
    if cycles = degree.find_all { |vertex, edges| edges.length > 0 } and cycles.length > 0
      message = cycles.collect { |vertex, edges| edges.collect { |e| e.to_s }.join(", ") }.join(", ")
      raise Puppet::Error, "Found dependency cycles in the following relationships: #{message}; try using the '--graph' option and open the '.dot' files in OmniGraffle or GraphViz"
    end

    result
  end

  # Add a new vertex to the graph.
  def add_vertex(vertex)
    @reversal = nil
    return false if vertex?(vertex)
    setup_vertex(vertex)
    true # don't return the VertexWrapper instance.
  end

  # Remove a vertex from the graph.
  def remove_vertex!(vertex)
    return nil unless vertex?(vertex)
    @vertices[vertex].edges.each { |edge| remove_edge!(edge) }
    @edges -= @vertices[vertex].edges
    @vertices[vertex].clear
    @vertices.delete(vertex)
  end

  # Test whether a given vertex is in the graph.
  def vertex?(vertex)
    @vertices.include?(vertex)
  end

  # Return a list of all vertices.
  def vertices
    @vertices.keys
  end

  # Add a new edge.  The graph user has to create the edge instance,
  # since they have to specify what kind of edge it is.
  def add_edge(source, target = nil, label = nil)
    @reversal = nil
    if target
      edge = Puppet::Relationship.new(source, target, label)
    else
      edge = source
    end
    [edge.source, edge.target].each { |vertex| setup_vertex(vertex) unless vertex?(vertex) }
    @vertices[edge.source].add_edge :out, edge
    @vertices[edge.target].add_edge :in, edge
    @edges << edge
    true
  end

  # Find a matching edge.  Note that this only finds the first edge,
  # not all of them or whatever.
  def edge(source, target)
    @edges.each_with_index { |test_edge, index| return test_edge if test_edge.source == source and test_edge.target == target }
  end

  def edge_label(source, target)
    return nil unless edge = edge(source, target)
    edge.label
  end

  # Is there an edge between the two vertices?
  def edge?(source, target)
    return false unless vertex?(source) and vertex?(target)

    @vertices[source].has_edge?(:out, target)
  end

  def edges
    @edges.dup
  end

  # Remove an edge from our graph.
  def remove_edge!(edge)
    @vertices[edge.source].remove_edge(:out, edge)
    @vertices[edge.target].remove_edge(:in, edge)

    @edges.delete(edge)
    nil
  end

  # Find adjacent edges.
  def adjacent(vertex, options = {})
    return [] unless wrapper = @vertices[vertex]
    wrapper.adjacent(options)
  end

  private

  # An internal method that skips the validation, so we don't have
  # duplicate validation calls.
  def setup_vertex(vertex)
    @vertices[vertex] = VertexWrapper.new(vertex)
  end

  public

#    # For some reason, unconnected vertices do not show up in
#    # this graph.
#    def to_jpg(path, name)
#        gv = vertices
#        Dir.chdir(path) do
#            induced_subgraph(gv).write_to_graphic_file('jpg', name)
#        end
#    end

  # Take container information from another graph and use it
  # to replace any container vertices with their respective leaves.
  # This creates direct relationships where there were previously
  # indirect relationships through the containers.
  def splice!(other, type)
    # We have to get the container list via a topological sort on the
    # configuration graph, because otherwise containers that contain
    # other containers will add those containers back into the
    # graph.  We could get a similar affect by only setting relationships
    # to container leaves, but that would result in many more
    # relationships.
    stage_class = Puppet::Type.type(:stage)
    whit_class  = Puppet::Type.type(:whit)
    containers = other.topsort.find_all { |v| (v.is_a?(type) or v.is_a?(stage_class)) and vertex?(v) }
    containers.each do |container|
      # Get the list of children from the other graph.
      children = other.adjacent(container, :direction => :out)

      # MQR TODO: Luke suggests that it should be possible to refactor the system so that
      #           container nodes are retained, thus obviating the need for the whit. 
      children = [whit_class.new(:name => container.name, :catalog => other)] if children.empty?

      # First create new edges for each of the :in edges
      [:in, :out].each do |dir|
        edges = adjacent(container, :direction => dir, :type => :edges)
        edges.each do |edge|
          children.each do |child|
            if dir == :in
              s = edge.source
              t = child
            else
              s = child
              t = edge.target
            end

            add_edge(s, t, edge.label)
          end

          # Now get rid of the edge, so remove_vertex! works correctly.
          remove_edge!(edge)
        end
      end
      remove_vertex!(container)
    end
  end

  # Just walk the tree and pass each edge.
  def walk(source, direction)
    # Use an iterative, breadth-first traversal of the graph. One could do
    # this recursively, but Ruby's slow function calls and even slower
    # recursion make the shorter, recursive algorithm cost-prohibitive.
    stack = [source]
    seen = Set.new
    until stack.empty?
      node = stack.shift
      next if seen.member? node
      connected = adjacent(node, :direction => direction)
      connected.each do |target|
        yield node, target
      end
      stack.concat(connected)
      seen << node
    end
  end

  # A different way of walking a tree, and a much faster way than the
  # one that comes with GRATR.
  def tree_from_vertex(start, direction = :out)
    predecessor={}
    walk(start, direction) do |parent, child|
      predecessor[child] = parent
    end
    predecessor
  end

  # LAK:FIXME This is just a paste of the GRATR code with slight modifications.

  # Return a DOT::DOTDigraph for directed graphs or a DOT::DOTSubgraph for an
  # undirected Graph.  _params_ can contain any graph property specified in
  # rdot.rb. If an edge or vertex label is a kind of Hash then the keys
  # which match +dot+ properties will be used as well.
  def to_dot_graph (params = {})
    params['name'] ||= self.class.name.gsub(/:/,'_')
    fontsize   = params['fontsize'] ? params['fontsize'] : '8'
    graph      = (directed? ? DOT::DOTDigraph : DOT::DOTSubgraph).new(params)
    edge_klass = directed? ? DOT::DOTDirectedEdge : DOT::DOTEdge
    vertices.each do |v|
      name = v.to_s
      params = {'name'     => '"'+name+'"',
        'fontsize' => fontsize,
        'label'    => name}
      v_label = v.to_s
      params.merge!(v_label) if v_label and v_label.kind_of? Hash
      graph << DOT::DOTNode.new(params)
    end
    edges.each do |e|
      params = {'from'     => '"'+ e.source.to_s + '"',
        'to'       => '"'+ e.target.to_s + '"',
        'fontsize' => fontsize }
      e_label = e.to_s
      params.merge!(e_label) if e_label and e_label.kind_of? Hash
      graph << edge_klass.new(params)
    end
    graph
  end

  # Output the dot format as a string
  def to_dot (params={}) to_dot_graph(params).to_s; end

  # Call +dotty+ for the graph which is written to the file 'graph.dot'
  # in the # current directory.
  def dotty (params = {}, dotfile = 'graph.dot')
    File.open(dotfile, 'w') {|f| f << to_dot(params) }
    system('dotty', dotfile)
  end

  # Use +dot+ to create a graphical representation of the graph.  Returns the
  # filename of the graphics file.
  def write_to_graphic_file (fmt='png', dotfile='graph')
    src = dotfile + '.dot'
    dot = dotfile + '.' + fmt

    File.open(src, 'w') {|f| f << self.to_dot << "\n"}

    system( "dot -T#{fmt} #{src} -o #{dot}" )
    dot
  end

  # Produce the graph files if requested.
  def write_graph(name)
    return unless Puppet[:graph]

    Puppet.settings.use(:graphing)

    file = File.join(Puppet[:graphdir], "#{name}.dot")
    File.open(file, "w") { |f|
      f.puts to_dot("name" => name.to_s.capitalize)
    }
  end
end