# Copyright 2007 Jeff Mesnil (http://jmesnil.net)
require 'java'
 
class String
  # Transform a CamelCase String to a snake_case String.
  #--
  # Code has been taken from ActiveRecord
  def snake_case
    self.to_s.gsub(/::/, '/').
    gsub(/([A-Z]+)([A-Z][a-z])/,'\1_\2').
    gsub(/([a-z\d])([A-Z])/,'\1_\2').
    tr("-", "_").
    downcase
  end
end
 
module JMX
  require 'dynamic_mbean'
  require 'open_data_helper'
  require 'objectname_helper'
  require 'jruby'
 
  class MBeanServerConnectionProxy
    attr_reader :connector
 
    # Adds a connector attribute to Java's native MBeanServerConnection class.
    #
    # The connector attribute can be used to manage the connection (e.g, to close it).
    # Why this isn't included in the native MBeanServerConnection class is beyond me.
    #
    # connector:: JMXConnector instance as returned by JMXConnectorFactory.connect.
    def initialize(connector)
      @connector = connector
      @connection = connector.getMBeanServerConnection
    end
 
    # Close the connection (an unfortunate omission from the MBeanServerConnection class, imho)
    def close
      @connector.close
    end
 
    # Forward all other method messages to the underlying MBeanServerConnection instance.
    def method_missing(method, *args, &block)
      @connection.send method, *args, &block
    end
  end
 
  class MBean
    include_class 'java.util.HashMap'
    include_class 'javax.naming.Context'
    include_class 'javax.management.Attribute'
    include_class 'javax.management.ObjectName'
    include_class 'javax.management.remote.JMXConnector'
    include_class 'javax.management.remote.JMXConnectorFactory'
    include_class 'javax.management.remote.JMXServiceURL'
    JThread = java.lang.Thread
 
    attr_reader :object_name, :operations, :attributes, :connection
 
    # Creates a new MBean.
    #
    # object_name:: a string corresponding to a valid ObjectName
    # connection::  a connection to a MBean server. If none is passed,
    #               use the global connection created by 
    #               MBean.establish_connection
    def initialize(object_name, connection=nil)
      @connection = connection || @@connection
      @object_name = object_name
      info = @connection.getMBeanInfo @object_name
      @attributes = Hash.new
      info.attributes.each do | mbean_attr |
        @attributes[mbean_attr.name.snake_case] = mbean_attr.name
        self.class.instance_eval do 
          define_method mbean_attr.name.snake_case do
            @connection.getAttribute @object_name, "#{mbean_attr.name}"
          end
        end
        if mbean_attr.isWritable
          self.class.instance_eval do
            define_method "#{mbean_attr.name.snake_case}=" do |value| 
              attr = Attribute.new mbean_attr.name, value
              @connection.setAttribute @object_name, attr
            end
          end
        end
      end
      @operations = Hash.new
      info.operations.each do |mbean_op|
        param_types = mbean_op.signature.map {|param| param.type}
        @operations[mbean_op.name.snake_case] = [mbean_op.name, param_types]
      end
    end
 
    def method_missing(method, *args, &block) #:nodoc:
      if @operations.keys.include?(method.to_s)
        op_name, param_types = @operations[method.to_s]
        @connection.invoke @object_name,
        op_name,
        args.to_java(:Object),
        param_types.to_java(:String)
      else
        super
      end
    end
 
    @@connection = nil
 
    # establish a connection to a remote MBean server which will
    # be used by all subsequent MBeans.
    # 
    # See MBean.create_connection for a list of the keys that are
    # accepted in arguments.
    #
    # Examples
    # 
    #   JMX::MBean.establish_connection :port => "node23", :port => 1090
    #   JMX::MBean.establish_connection :port => "node23", :username => "jeff", :password => "secret"
    def self.establish_connection(args={})
      @@connection ||= create_connection args
    end
 
    def self.remove_connection(args={})
      if @@connection
        @@connection.close rescue nil
      end
      @@connection = nil
    end
 
    def self.connection(args={})
      if args.has_key? :host or args.has_key? :port
        return create_connection(args)
      else
        @@connection ||= MBean.establish_connection(args)
      end
    end
 
    # Create a connection to a remote MBean server.
    # 
    # The args accepts the following keys:
    #
    # [:host]             the host of the MBean server (defaults to "localhost")
    #
    # [:port]             the port of the MBean server (defaults to 3000)
    #
    # [:url]              the url of the MBean server.
    #                     No default.
    #                     if the url is specified, the host & port parameters are
    #                     not taken into account
    #
    # [:username]         the name of the user (if the MBean server requires authentication).
    #                     No default
    #
    # [:password]         the password of the user (if the MBean server requires authentication).
    #                     No default
    #
    # [:credentials]      custom credentials (if the MBean server requires authentication).
    #                     No default. It has precedence over :username and :password (i.e. if
    #                     :credentials is specified, :username & :password are ignored)   
    #
    # [:provider_package] use to fill the JMXConnectorFactory::PROTOCOL_PROVIDER_PACKAGES.
    #                     No default
    #
    def self.create_connection(args={})
      host= args[:host] || "localhost"
      port = args[:port] || 3000
      username = args[:username]
      password = args[:password]
      credentials = args[:credentials]
      provider_package = args[:provider_package]
      
      # host & port are not taken into account if url is set (see issue #7)
      standard_url = "service:jmx:rmi:///jndi/rmi://#{host}:#{port}/jmxrmi"
      url = args[:url] || standard_url
      
      unless credentials
        if !username.nil? and username.length > 0
          user_password_credentials = [username, password]
          credentials = user_password_credentials.to_java(:String)
        end
      end
      
      env = HashMap.new
      env.put(JMXConnector::CREDENTIALS, credentials) if credentials
      # only fill the Context and JMXConnectorFactory properties if provider_package is set
      if provider_package
        env.put(Context::SECURITY_PRINCIPAL, username) if username
        env.put(Context::SECURITY_CREDENTIALS, password) if password
        env.put(JMXConnectorFactory::PROTOCOL_PROVIDER_PACKAGES, provider_package)
      end
 
      # the context class loader is set to JRuby's classloader when
      # creating the JMX Connection so that classes loaded using 
      # JRuby "require" (and not from its classpath) can also be 
      # accessed (see issue #6)
      begin
        context_class_loader = JThread.current_thread.context_class_loader
        JThread.current_thread.context_class_loader = JRuby.runtime.getJRubyClassLoader
        
        connector = JMXConnectorFactory::connect JMXServiceURL.new(url), env
        MBeanServerConnectionProxy.new connector
      ensure
        # ... and we reset the previous context class loader
        JThread.current_thread.context_class_loader = context_class_loader
      end
    end
 
    # Returns an array of MBeans corresponding to all the MBeans
    # registered for the ObjectName passed in parameter (which may be 
    # a pattern).
    #
    # The args accepts the same keys than #create_connection and an 
    # additional one:
    #
    # [:connection] a MBean server connection (as returned by #create_connection)
    #               No default. It has precedence over :host and :port (i.e if
    #               :connection is specified, :host and :port are ignored)
    #
    def self.find_all_by_name(name, args={})
      object_name = ObjectName.new(name)
      connection = args[:connection] || MBean.connection(args)
      object_names = connection.queryNames(object_name, nil)
      object_names.map { |on| MBean.new(on, connection) }
    end
 
    # Same as #find_all_by_name but the ObjectName passed in parameter
    # can not be a pattern.
    # Only one single MBean is returned.
    def self.find_by_name(name, args={})
      connection = args[:connection] || MBean.connection(args)
      MBean.new ObjectName.new(name), connection
    end
 
    def self.pretty_print (object_name, args={})
      connection = args[:connection] || MBean.connection(args)
      info = connection.getMBeanInfo ObjectName.new(object_name)
      puts "object_name: #{object_name}"
      puts "class: #{info.class_name}"
      puts "description: #{info.description}"
      puts "operations:"
      info.operations.each do | op |
        puts "  #{op.name}"
        op.signature.each do | param |
          puts "    #{param.name} (#{param.type} #{param.description})"
        end
        puts "    ----"
        puts "    description: #{op.description}"
        puts "    return_type: #{op.return_type}"
        puts "    impact: #{op.impact}"
      end
      puts "attributes:"
      info.attributes.each do | attr |
        puts "  #{attr.name}"
        puts "    description: #{attr.description}"
        puts "    type: #{attr.type}"
        puts "    readable: #{attr.readable}"
        puts "    writable: #{attr.writable}"
        puts "    is: #{attr.is}"
      end
    end
  end
end