#!/usr/bin/ruby
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

require 'json'
require 'emr/common'

def run(cmd)
    if ! system(cmd) then
      raise "Command failed: #{cmd}"
    end
end

def sudo(cmd)
   run("sudo #{cmd}")
end

def println(*args)
    print *args
    puts
end

job_flow = Emr::JsonInfoFile.new('job-flow')
instance_info = Emr::JsonInfoFile.new('instance')
is_master = instance_info['isMaster'].to_s == 'true'
$zookeeper = job_flow['masterPrivateDnsName']

def cleanUp(targetDir,runDir,logDir)
    println "Cleaning up stale directories.. "
    println "\t#{targetDir}/*drill*"
    sudo "rm -rf #{targetDir}/*drill*"
    println "\t/home/hadoop/drill"
    sudo "rm -rf /home/hadoop/drill"
    println "\t#{runDir} #{logDir}"
    sudo "rm -rf #{runDir} #{logDir}"
end

def writeDrillConfs(runDir,logDir)
    zookeeperConnect=$zookeeper
    File.open('/tmp/drill-override.conf', 'w') do |f|
    f.write(<<EOF
drill.exec: {
  cluster-id: "EMR-Drill"
  rpc: {
        user.port : 31010,
        bit.port : 32011
  },
  operator: {
    packages += "org.apache.drill.exec.physical.config"
  },
  optimizer: {
    implementation: "org.apache.drill.exec.opt.IdentityOptimizer"
  },
  storage: {
        packages += "org.apache.drill.exec.store"
  }
  metrics : {
        context: "drillbit"
  },
  zk: {
        connect: "#{zookeeperConnect}",
        root: "drill",
        refresh: 500,
        timeout: 5000,
        retry: {
          count: 7200,
          delay: 500
        }
  }

  network: {
    start: 35000
  }
}
EOF
)
end

    sudo "mv /tmp/drill-override.conf /home/hadoop/drill/conf/drill-override.conf && chmod 644 /home/hadoop/drill/conf/drill-override.conf"

    File.open('/tmp/drillbit', 'w') do |f| 
    f.write(<<EOF
/home/hadoop/drill/bin/drillbit.sh $@
EOF
)
end
    sudo "mv /tmp/drillbit /etc/init.d/drillbit && chmod a+x /etc/init.d/drillbit"

    File.open('/tmp/drill-init-config.sh', 'w') do |f|
    f.write(<<EOF
export DRILL_LOG_DIR=/mnt/var/log/drill
export DRILL_PID_DIR=/mnt/var/run/drill
export DRILL_MAX_DIRECT_MEMORY=4096M
EOF
)
end
    sudo "cat /tmp/drill-init-config.sh >> /home/hadoop/drill/bin/drill-config.sh && rm /tmp/drill-init-config.sh"

    File.open('/tmp/drill.conf', 'w') do |f| 
    f.write(<<EOF
[

{
  "name": "run-drill-zookeeper",
  "type": "file",
  "file": "#{runDir}/run-zookeeper",
  "pattern": "1"
},

{
  "name": "run-drill",
  "type": "file",
  "file": "#{runDir}/run-drill",
  "pattern": "1"
},

{
  "name": "drill-bit",
  "type": "process",
  "start": "/etc/init.d/drillbit start",
  "stop": "/etc/init.d/drillbit stop",
  "pid-file": "#{runDir}/drillbit.pid",
  "pattern": "org.apache.drill.exec.server.Drillbit",
  "depends": ["run-drill", "drill-zookeeper"]
},

{
  "name": "drill-zookeeper",
  "type": "process",
  "start": "/home/hadoop/drill/zookeeper/bin/zkServer.sh start",
  "stop": "/home/hadoop/drill/zookeeper/bin/zkServer.sh stop",
  "pid-file": "#{runDir}/zookeeper/zookeeper_server.pid",
  "pattern": "org.apache.zookeeper.server.quorum.QuorumPeerMain",
  "depends": ["run-drill-zookeeper"]
}

]
EOF
)
end

    sudo "mv /tmp/drill.conf /etc/service-nanny/drill.conf"
    sudo "echo '1' >/mnt/var/run/drill/run-zookeeper"
    sudo "echo '1' >/mnt/var/run/drill/run-drill"

end

def installZookeeper(targetDir, runDir, logDir)
    zkTarGZUrl="http://beta.elasticmapreduce.s3.amazonaws.com/bootstrap-actions/zookeeper-standalone"
    zkTarGZ="zookeeper-3.4.5.tar.gz"

    println "Going to download Zookeeper distribution #{zkTarGZ} from #{zkTarGZUrl}"
    sudo "curl -L --silent --show-error --fail --connect-timeout 60 --max-time 720 --retry 5 -O  #{zkTarGZUrl}/#{zkTarGZ}"

    println "Installing zookeeper in #{targetDir}"
    sudo "tar xzf #{zkTarGZ} && rm -f #{zkTarGZ}"
    sudo "mv *zookeeper* #{targetDir}/"
    sudo "ln -s `ls -d #{targetDir}/*zookeeper*` /home/hadoop/drill/zookeeper"

    File.open('/tmp/zoo.cfg', 'w') do |f|
    f.write(<<EOF
maxClientCnxns=50

# The number of milliseconds of each tick
tickTime=2000

# The number of ticks that the initial 
# synchronization phase can take
initLimit=10

# The number of ticks that can pass between 
# sending a request and getting an acknowledgement
syncLimit=5

# the directory where the snapshot is stored.
dataDir=#{runDir}

# the directory where the snapshot is stored.
dataLogDir=#{logDir}

# the port at which the clients will connect
clientPort=2181
EOF
)
end
    sudo "mv /tmp/zoo.cfg /home/hadoop/drill/zookeeper/conf/ && chmod 644 /home/hadoop/drill/zookeeper/conf/zoo.cfg"

end

def installDrill(targetDir, runDir, logDir)
    drillTarGZUrl="http://getdrill.org/drill/download"
    drillTarGZ="apache-drill-1.0.0.tar.gz"
    drillRoot="apache-drill-1.0.0"

    println "Going to download Apache Drill distribution #{drillTarGZ} from #{drillTarGZUrl}"
    sudo "curl -L --silent --show-error --fail --connect-timeout 60 --max-time 720 --retry 5 -O  #{drillTarGZUrl}/#{drillTarGZ}"

    println "Installing Apache Drill..."
    sudo "tar xzf #{drillTarGZ} && rm -f #{drillTarGZ}"
    sudo "mv #{drillRoot} #{targetDir}/ && ln -s #{targetDir}/#{drillRoot} /home/hadoop/drill"

    sudo "mkdir -p #{runDir}/zookeeper"
    sudo "mkdir -p #{logDir}/zookeeper"
    println "Installing Zookeeper..."
    installZookeeper("/home/hadoop/drill", "#{runDir}/zookeeper", "#{logDir}/zookeeper")

    println "Configuring Apache Drill..."
    sudo "chown -R hadoop.hadoop /home/hadoop/.versions/*drill*"
    sudo "mkdir -p #{runDir} && mkdir -p #{logDir}"
    sudo "chown -R hadoop.hadoop #{runDir}"
    sudo "chown -R hadoop.hadoop #{logDir}"

    writeDrillConfs($runDir,$logDir)

    sudo "chown -R hadoop.hadoop #{runDir}"
    sudo "chown -R hadoop.hadoop #{logDir}"

    sudo "/etc/init.d/service-nanny start"
    sudo "/etc/init.d/service-nanny restart"
end

$targetDir="/home/hadoop/.versions"
$runDir="/mnt/var/run/drill"
$logDir="/mnt/var/log/drill"

cleanUp($targetDir, $runDir, $logDir)
installDrill($targetDir, $runDir, $logDir)