#!/usr/bin/env ruby
# grid.rb
# Grid Helper

require 'ap'
require 'watirgrid'
require 'json'
require 'rest_client'
require 'highline/import'

class Grid

  def initialize(params = {})
    @log = Logger.new(STDOUT, 'daily')
    @log.level = params[:loglevel] || Logger::INFO
    @exclusive = params[:exclusive]
    @exclusive ? params[:take_all] = true : params[:read_all] = true
    @webdriver_browser_type = params[:browser].to_sym if params[:browser]
    @grid = Watir::Grid.new(params)
    @grid.start(params)
    @providers = @grid.browsers
    @browsers = []
  end

  def size
    @grid.size
  end

  def browsers
    @browsers
  end

  def providers
    @providers
  end

  def setup
    @grid.browsers.each_with_index do |browser, index|
      sleep 0.15
      @browsers[index] ||= browser[:object].new_browser(@webdriver_browser_type)
    end
  end

  def teardown
    @browsers.each do |browser|
      browser.close
    end
    @grid.release_tuples if @exclusive
  end

  def iterate &block
    threads = []
    @browsers.each do |browser|
      threads << Thread.new do
        yield browser
      end
    end
    threads.each {|thread| thread.join}
  end

  def iterate_with_index &block
    threads = []
    @browsers.each_with_index do |browser, index|
      threads << Thread.new do
        yield browser, index
      end
    end
    threads.each {|thread| thread.join}
  end

  def self.control(params = {}, &block)
    @exclusive = params[:exclusive]
    @exclusive ? params[:take_all] = true : params[:read_all] = true
    @webdriver_browser_type = params[:browser].to_sym if params[:browser]
    @grid = Watir::Grid.new(params)
    @grid.start(params)
    @browsers = []
    threads = []
    @grid.browsers.each_with_index do |browser, index|
      sleep rampup(params)
      threads << Thread.new(@webdriver_browser_type) do |webdriver_browser_type|
        @browser = browser[:object].new_browser(webdriver_browser_type)
        yield @browser, index
      end
    end
    threads.each {|thread| thread.join}
    @grid.release_tuples if @exclusive
  end

  private

  def self.rampup(params = {})
    if params[:rampup]
      params[:rampup] / @grid.size
    else
      0.5
    end
  end

  class Helper
    def initialize(params = {})
      @grid_id  = params[:grid_id]
      @api      = params[:api]
      @uri = params[:uri] || "http://gridin.it/api/v0"
      logfile = STDOUT
      @log  = Logger.new(logfile, 'daily')
      @log.level = Logger::DEBUG
      @log.datetime_format = "%Y-%m-%d %H:%M:%S "
      if home
        if File.file? ENV[home]+'/.gridrc'
          @token = open(ENV[home]+'/.gridrc') { |f| f.read }.split(':')[1]
        end
      end
      @token = get_token(params) unless @token
      @log.debug("Current token: #{@token}")
    end

    def get_password(prompt='Password: ', mask='*')
      ask(prompt) { |q| q.echo = mask }
    end

    def get_token(params={})
      if params[:token]
        params[:token]
      else
        HighLine.track_eof = false
        @email    = params[:email] || ask('Email: ')
        @password = params[:password] ||  get_password
        response = RestClient.post "#{@uri}/token.json" , {
          :email => @email, :password => @password
        }
        (JSON.parse(response))["token"]
      end
    end

    def home
      homes = ["HOME", "HOMEPATH"]
      homes.detect {|h| ENV[h] != nil}
    end

    def token
      @token = get_token
      open(ENV[home]+'/.gridrc', 'w') { |f| f << "TOKEN:#{@token}" } if home and @token
      @log.debug("New token    : #{@token}")
    end

    def list
      @log.debug("Listing grids ...")
      begin
        response = RestClient.get "#{@uri}/grid.json" , {
          :token => @token
        }
        ap JSON.parse(response)
      rescue => e
        @log.error("#{e}")
      end
    end

    def show
      @grid_id = ask("Grid ID: ") unless @grid_id
      @log.debug("Listing grid ...")
      begin
        response = RestClient.get "#{@uri}/grid/#{@grid_id}.json" , {
          :token => @token
        }
        ap JSON.parse(response)
      rescue => e
        @log.error("#{e}")
      end
    end

    def create_grid(name, description)
      begin
        @log.debug("Creating grid ...")
        response = RestClient.post "#{@uri}/grid.json" , {
          :token => @token,
          :name  => name,
          :description => description
        }
        (JSON.parse(response)).first[1]["id"]
      rescue => e
        @log.error("#{e}")
      end
    end

    def create_intranode(grid_id, browser_type, location, nodes)
      begin
        @log.debug("Creating intranode ...")
        response = RestClient.post "#{@uri}/grid/#{grid_id.to_s}/intranode.json", {
          :token => @token,
          :browser_type => browser_type,
          :location => location,
          :nodes => nodes
        }
      rescue => e
        @log.error("#{e}")
      end
    end

    def delete_intranode
      begin
        @log.debug("Deleting intranode ...")
        response = RestClient.delete "#{@uri}/grid/intranode/#{@intranode_id}.json", {
          :token => @token
        }
      rescue => e
        @log.error("#{e}")
      end
    end

    def create
      name = ask("Grid name: ") { |q| q.default = "Grid" }
      description = ask("Grid description: ") { |q| q.default = "Grid created #{Time.now}" }
      nodes = ask("Number of nodes: ") { |q| q.default = 5 }
      location = ask("Location of nodes: ") { |q| q.default = "us-east" }
      browser_type = ask("Node type: ") { |q| q.default = "webdriver" } unless @api
      @grid_id = create_grid(name, description)
      create_intranode(@grid_id, browser_type, location, nodes)
      show
    end

    def add
      @grid_id = ask("Grid ID: ") unless @grid_id
      nodes = ask("Number of nodes: ") { |q| q.default = 5 }
      location = ask("Location of nodes: ") { |q| q.default = "us-east" }
      browser_type = ask("Node type: ") { |q| q.default = "webdriver" } unless @api
      create_intranode(@grid_id, browser_type, location, nodes)
      show
    end

    def delete
      @grid_id = ask("Grid ID: ") unless @grid_id
      @intranode_id= ask("Intranode ID: ") unless @intranode_id
      delete_intranode
      show
    end

    def start
      @grid_id = ask("Grid ID: ") unless @grid_id
      @log.debug("Starting grid ...")
      begin
        response = RestClient.put "#{@uri}/grid/#{@grid_id.to_s}/start.json", {:token => @token}
      rescue => e
        @log.error("#{e}")
      end
      show
    end

    def stop
      @grid_id = ask("Grid ID: ") unless @grid_id
      @log.debug("Stopping grid ...")
      begin
        response = RestClient.put "#{@uri}/grid/#{@grid_id.to_s}/stop.json", {:token => @token}
      rescue => e
        @log.error("#{e}")
      end
      show
    end

    def restart
      @grid_id = ask("Grid ID: ") unless @grid_id
      @log.debug("Restarting grid ...")
      begin
        response = RestClient.put "#{@uri}/grid/#{@grid_id.to_s}/restart.json", {:token => @token}
      rescue => e
        @log.error("#{e}")
      end
      show
    end

    def status
      @grid_id = ask("Grid ID: ") unless @grid_id
      @log.debug("Updating status of grid ...")
      begin
        response = RestClient.put "#{@uri}/grid/#{@grid_id.to_s}/status.json", {:token => @token}
      rescue => e
        @log.error("#{e}")
      end
      show
    end

  end
end