#!/usr/bin/env ruby
#
#  console.rb: Console GUI component.
#
#  Copyright (C) 2011 VMware, Inc.
# 
#  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.

#  Apologies in advance for the state of this code, my Ruby/Tk-fu is really not
#  up to scratch. The liberal use of $global_variables is to get things working
#  correctly from Tk callbacks, if you know a better way to make this work
#  please let me know.

require 'tk'
require 'ffi-rzmq'
require 'thread'

#  Convenience method to raise exception on ffi-rzmq error
def raise_if_error(rc)
  unless ZMQ::Util.resultcode_ok?(rc)
    raise "ZMQ Error: #{ZMQ::Util.error_string}"
  end
end

#  Callable on any Tk widget, adds get/set methods for mytag and mytype,
#  allowing us to keep track of what is what when we get callback events.
def make_taggable(widget)
  widget.instance_eval do
    @mytag = nil
    @mytype = nil
    def mytag
      @mytag
    end
    def mytag=(tag)
      @mytag = tag
    end
    def mytype
      @mytype
    end
    def mytype=(type)
      @mytype = type
    end
  end
end

#  Build out the GUI.
def build_gui
  $root = TkRoot.new { title "Console" }
  $workspace = TkFrame.new($root) {
      relief 'sunken'
      borderwidth '2'
    }
  $workspace.grid :column => 0, :row => 0, :columnspan => 2, :sticky => 'nsew'
  $buttons = TkFrame.new($root)
  $buttons.grid :column => 0, :row => 1, :columnspan => 2, :sticky => 'nsew'
  $start = TkButton.new($buttons) {
      text 'Start'
      state :disabled
      command proc {
        host = $selected_item.mytag
        start_service_dialog(host)
      }
    }
  $start.pack :padx => 10, :side => :left
  $stop = TkButton.new($buttons) {
      text 'Stop'
      state :disabled
      command proc {
        service = $selected_item.mytag
        host = $selected_item.winfo_parent.mytag
        result = Tk::messageBox :type => :yesno,
          :message => "Are you sure you wish to stop the service " + \
            "'#{service}' on host '#{host}'?",
          :icon => :question,
          :title => "Stop service"
        stop_service(service, host) if result == 'yes'
      }
    }
  $stop.pack :padx => 10, :side => :left
  $restart = TkButton.new($buttons) {
      text 'Restart'
      state :disabled
      command proc {
        service = $selected_item.mytag
        host = $selected_item.winfo_parent.mytag
        restart_service_dialog(host, service)
      }
    }
  $restart.pack :padx => 10, :side => :left
  $dns = TkButton.new($buttons) {
      text 'Update DNS'
      command proc {
        update_dns_dialog()
      }
    }
  $dns.pack :padx => 10, :side => :left
  $exit = TkButton.new($buttons) {
      text 'Exit'
      command proc {
        result = Tk::messageBox :type => :yesno,
          :message => "Are you sure you want to exit?",
          :icon => :question,
          :title => "Confirm exit"
        exit if result == 'yes'
      }
    }
  $exit.pack :side => :left
  $console = TkListbox.new($root) { 
      height 10
      width 60
      yscrollcommand proc { |*args| $console_sb.set(*args) }
    }
  $console.grid :column => 0, :row => 2, :sticky => 'nwes'
  $console_sb = TkScrollbar.new($root) {
      orient :vertical
      command proc { |*args| $console.yview(*args) }
    }
  $console_sb.grid :column => 1, :row => 2, :sticky => 'ns'
  TkGrid.columnconfigure $root, 0, :weight => 1, :minsize => 600
  TkGrid.rowconfigure $root, 0, :weight => 1, :minsize => 400
end

#  DNS Update dialog GUI
def update_dns_dialog
  w = TkToplevel.new { title "Update DNS record" }
  w.wm_transient($root)
  l = TkLabel.new(w) { text "Record" }
  r = TkEntry.new(w) { width 15 }
  r.focus
  l2 = TkLabel.new(w) { text "Value" }
  v = TkEntry.new(w) { width 15 }
  ok = TkButton.new(w) {
    text "OK"
    command proc {
      if r.get.length > 0 && v.get.length > 0
        result = `./nsupdate-script #{r.get} #{v.get}`
        result.chomp
        console_add("Updating DNS record '#{r.get}' to value '#{v.get}': #{result}")
        w.destroy
      end
      }
    default :active
    }
  cancel = TkButton.new(w) { text "Cancel"; command proc { w.destroy } }
  w.bind('Return') { ok.invoke }
  w.bind('Escape') { cancel.invoke }
  l.grid :column => 0, :row => 0, :padx => 10, :pady => 10
  r.grid :column => 1, :row => 0, :padx => 10, :pady => 10
  l2.grid :column => 0, :row => 1, :padx => 10, :pady => 10
  v.grid :column => 1, :row => 1, :padx => 10, :pady => 10
  ok.grid :column => 0, :row => 2, :padx => 10, :pady => 10
  cancel.grid :column => 1, :row => 2, :padx => 10, :pady => 10
  w.grab
  w.tkwait
end

#  Start Service dialog GUI
def start_service_dialog(host)
  w = TkToplevel.new { title "Start service on #{host}" }
  w.wm_transient($root)
  l = TkLabel.new(w) { text "Command" }
  e = TkEntry.new(w) { width 40 }
  e.focus
  ok = TkButton.new(w) {
    text "OK"
    command proc {
      if e.get.length > 0
        start_service(host, e.get)
        w.destroy
      end
      }
    default :active
    }
  cancel = TkButton.new(w) { text "Cancel"; command proc { w.destroy } }
  w.bind('Return') { ok.invoke }
  w.bind('Escape') { cancel.invoke }
  l.grid :column => 0, :row => 0, :padx => 10, :pady => 10
  e.grid :column => 1, :row => 0, :padx => 10, :pady => 10
  ok.grid :column => 0, :row => 1, :padx => 10, :pady => 10
  cancel.grid :column => 1, :row => 1, :padx => 10, :pady => 10
  w.grab
  w.tkwait
end

#  Restart service dialog GUI
def restart_service_dialog(host, service)
  w = TkToplevel.new { title "Restart service #{service} on #{host}" }
  w.wm_transient($root)
  l = TkLabel.new(w) { text "Command" }
  e = TkEntry.new(w) { width 40 }
  e.set($hosts[host][:services][service][:command])
  e.focus
  ok = TkButton.new(w) {
    text "OK"
    command proc {
      if e.get.length > 0
        stop_service(service, host)
        start_service(host, e.get)
        w.destroy
      end
      }
    default :active
    }
  cancel = TkButton.new(w) { text "Cancel"; command proc { w.destroy } }
  w.bind('Return') { ok.invoke }
  w.bind('Escape') { cancel.invoke }
  l.grid :column => 0, :row => 0, :padx => 10, :pady => 10
  e.grid :column => 1, :row => 0, :padx => 10, :pady => 10
  ok.grid :column => 0, :row => 1, :padx => 10, :pady => 10
  cancel.grid :column => 1, :row => 1, :padx => 10, :pady => 10
  w.grab
  w.tkwait
end

#  Start service 'command' on 'host'
def start_service(host, command)
  console_add("Starting command '#{command}' on host #{host}");
  s = $ctx.socket(ZMQ::PUSH)
  rc = s.connect("tcp://#{host}:7772")
  raise_if_error(rc)
  rc = s.send_string("START #{command}")
  raise_if_error(rc)
  rc = s.close
  raise_if_error(rc)
end

#  Stop service 'service' on 'host'
def stop_service(service, host)
  console_add("Stopping service: #{service} on host: #{host}")
  s = $ctx.socket(ZMQ::PUSH)
  rc = s.connect("tcp://#{host}:7772")
  raise_if_error(rc)
  rc = s.send_string("STOP #{service}")
  raise_if_error(rc)
  rc = s.close
  raise_if_error(rc)
end

#  Log text to console widget
def console_add(msg)
  $console.insert 'end', Time.now.to_s + ' ' + msg
  $console.yview 'end'
end

#  Select an item in the console, highlights it and updates available
#  action buttons
def select_item(new)
  old = $selected_item
  if old
    target = old.labelwidget if old.mytype == :host
    target = old if old.mytype == :service
    target.background('#d9d9d9')
  end
  if new.mytype == :host
    new.labelwidget.background('green')
    $start.state(:normal)
    $stop.state(:disabled)
    $restart.state(:disabled)
  elsif new.mytype == :service
    new.background('green')
    $start.state(:disabled)
    $stop.state(:normal)
    $restart.state(:normal)
  end
  $selected_item = new
end

#  Run periodically from Tk Timer to pull interesting messages from notifier
#  and update the GUI accordingly.
def update
  begin
    msg = ''
    rc = 0
    loop do
      $semaphore.synchronize do
        rc = $notifier_sock.recv_string(msg, ZMQ::DONTWAIT)
      end
      if rc < 0
        break if ZMQ::Util.errno == ZMQ::EAGAIN
        raise_if_error(rc)
      end
      if msg =~ /^WATCHDOG (\w+) (\w+) ([\w\.]+) ?(.*)?$/
        verb = $1
        host = $2
        timestamp = $3.to_f
        command = $4
        if verb == 'ALIVE'
          if $hosts.has_key?(host)
            $hosts[host][:last_seen] = timestamp
          else
            label = TkLabel.new($workspace) { text host }
            make_taggable(label)
            label.mytag = host
            label.mytype = :host
            host_widget = TkLabelFrame.new($workspace) {
              labelwidget label
              width '4c'
              height '4c'
              relief 'raised'
              padx 10
              pady 10
            }
            make_taggable(host_widget)
            host_widget.mytag = host
            host_widget.mytype = :host
            host_widget.pack :side => :left,:anchor => 'nw',
              :padx => 10, :pady => 10
            host_widget.cursor("hand2")
            label.cursor("hand2")
            label.bind("1", proc { |event|
              host = event.widget.mytag
              real_widget = $hosts[host][:widget]
              select_item(real_widget)
              console_add("Selected host: #{host}")
              })
            host_widget.bind("1", proc { |event|
              select_item(event.widget)
              console_add("Selected host: #{event.widget.mytag}")
              })
            $hosts[host] = {}
            $hosts[host][:widget] = host_widget
            $hosts[host][:services] = {}
            $hosts[host][:last_seen] = timestamp
            console_add("Registered host: #{host}")
          end
        elsif verb == 'RUNNING'
          service = command.split()[0]
          svcs = $hosts[host][:services]
          if svcs.has_key?(service)
            svcs[service][:last_seen] = timestamp
          else
            svcs[service] = {}
            svcs[service][:command] = command
            svcs[service][:last_seen] = timestamp
            host_widget = $hosts[host][:widget]
            service_widget = TkLabel.new(host_widget) { 
              text service
              padx '0.25c'
              pady '0.25c'
              }
            make_taggable(service_widget)
            service_widget.mytype = :service
            service_widget.mytag = service
            service_widget.pack { fill 'both' }
            service_widget.cursor("hand2")
            service_widget.bind("1", proc { |event|
              select_item(event.widget)
              console_add("Selected service: #{event.widget.mytag}")
              })
            svcs[service][:widget] = service_widget
            console_add("Registered service: #{service} on host: #{host}")
          end
        end
      elsif msg =~ /^VALUE (\w+) (\w+) (.*)$/
        host = $1
        service = $2
        value = $3
        if $hosts.has_key?(host) && $hosts[host][:services].has_key?(service)
          widget = $hosts[host][:services][service][:widget]
          widget.text(widget.mytag + "\n" + value)
        end
      end
    end
  rescue => exception
    print exception, "\n"
    print exception.backtrace.join("\n\t from "), "\n"
    Process.abort
  end
end

#  Run periodically from Tk Timer to reap dead hosts/services from GUI.
def reaper
  begin
    current_time = Time.now.to_f
    $hosts.each do |key, host|
      if current_time - host[:last_seen] > 2
        host[:services].each do |key, service|
          if $selected_item == service[:widget]
            $start.state(:disabled)
            $stop.state(:disabled)
            $restart.state(:disabled)
            $selected_item = nil
          end
          service[:widget].destroy
          host[:services].delete(key)
        end
        if $selected_item == host[:widget]
          $start.state(:disabled)
          $stop.state(:disabled)
          $restart.state(:disabled)
          $selected_item = nil
        end
        host[:widget].destroy
        $hosts.delete(key)
        console_add("Removed host: #{key}")
      else
        host[:services].each do |key, service|
          if current_time - service[:last_seen] > 2
            if $selected_item == service[:widget]
              $start.state(:disabled)
              $stop.state(:disabled)
              $restart.state(:disabled)
              $selected_item = nil
            end
            service[:widget].destroy
            host[:services].delete(key)
            console_add("Removed service: #{key}")
          end
        end
      end
    end
  rescue => exception
    print exception, "\n"
    print exception.backtrace.join("\n\t from "), "\n"
    Process.abort
  end
end

#
#  Main program
#

#  Globals used by GUI
#  ZMQ context
$ctx = ZMQ::Context.new()
#  Currently selected item (host/service)
$selected_item = nil
#  Map of hosts and their services
$hosts = {}

#  Parse arguments
unless ARGV[0]
  print "usage: console.rb <notifier-host>\n"
  exit 1
end
notifier_ep = "tcp://#{ARGV[0]}:7771"

#  Create and connect to notifier socket. Synchronized to allow migration to
#  GUI callback thread.
$semaphore = Mutex.new
$semaphore.synchronize do
  $notifier_sock = $ctx.socket(ZMQ::SUB)
  rc = $notifier_sock.setsockopt(ZMQ::SUBSCRIBE, '')
  raise_if_error(rc)
  rc = $notifier_sock.connect(notifier_ep)
  raise_if_error(rc)
end

#  Build and update GUI immediately on startup
build_gui
update

#  Setup timers and go into Tk mainloop
#  Update GUI/reap dead objects every 500ms
TkTimer.new(500, -1, proc { update }).start
TkTimer.new(500, -1, proc { reaper }).start
Tk.mainloop