#!/usr/bin/env ruby

config = {
  data_dir: nil,
  command: nil,
  binary: nil,
}

def usage(error=nil)
  puts "usage: #{File.basename($0)} -D <data dir> [-P <postgresql bin path>] <command>"
  puts "commands:"
  puts "  setup                 setup <data dir>"
  puts "  pgpool                start pgpool as a daemon process"
  puts "  pgpool stop           stop  pgpool daemon process"
  puts "  pgpool -n             start pgpool as a foreground process"
  puts "  pg_ctl start          start postgres server as a daemon process"
  puts "  pg_ctl stop           stop  postgres server daemon process"
  puts "  postgres              start postgres server as a foreground process"
  puts "  passwd <USER NAME>    add new md5 password entry for an user"
  puts "  show_init_sql         display statements to initialize a new PostgreSQL database"
  puts ""
  puts "error: #{error}" if error
  exit 0
end

config_key = false

arg_offset = 0
ARGV.each_with_index do |a,i|
  arg_offset = i + 1

  if config_key
    config[config_key] = a
    config_key = nil
    next
  end

  case a
  when "-D"
    config_key = :data_dir
  when "-P"
    config_key = :bin_path
  when "-h", "-?", "--help"
    config[:command] = :help
  when "setup"
    config[:command] = :setup
  when "pgpool"
    config[:command] = :pgpool
    break
  when "pg_ctl"
    config[:command] = :pg_ctl
    break
  when "postgres"
    config[:command] = :postgres
    break
  when "passwd"
    config[:command] = :passwd
    break
  when "show_init_sql"
    config[:command] = :show_init_sql
    break
  end
end

usage nil if config_key
usage nil unless config[:command]

args = ARGV[arg_offset..-1]

setup_params = {
 port: 9900,
 backend_port: 6432,
 pgpool_pid_file: nil,
 pgpool_status_dir: nil,
 unix_socket_directory: nil,
}

ENV['PATH'] = "#{ENV['PATH']}:#{config[:bin_path]}"

case config[:command]
when :help
  usage nil

when :setup
  unless config[:data_dir]
    usage "-D option is required"
  end
  unless args.empty?
    usage nil
  end

  require "erb"
  require "fileutils"
  require "shellwords"

  def setup_cmd(cmdline)
    puts "setup> #{cmdline}"
    system cmdline
    unless $?.success?
      puts "******************"
      puts "** setup failed **"
      puts "******************"
      puts ""
      puts "  * -D directory is inconsistent."
      puts ""
      puts "  * If you're running command on Mac OS X, you might need to run following commands:"
      puts ""
      puts "    $ sudo sysctl -w kern.sysv.shmmax=1073741824"
      puts "    $ sudo sysctl -w kern.sysv.shmall=1073741824"
      puts ""
      puts "  * Please delete the directory (rm -rf) first before retrying."
      puts ""
      exit 1
    end
  end

  def se(raw)
    Shellwords.escape(raw)
  end

  `which initdb && which postgres`
  unless $?.success?
    puts ""
    puts "Can't find initdb or postgres command. Check following items:"
    puts ""
    puts "  * You need to install PostgreSQL server. Use apt or yum to install postgresql package."
    puts "  * You need to set -P option to specify path to the commands. Example:"
    puts ""
    puts "    $ prestogres -D data_dir -P /usr/lib/postgresql/9.1/bin setup"
    puts ""
    exit 1
  end

  # src files
  config_src_dir = File.join(File.dirname(__FILE__), "..", "config")
  postgresql_conf_path = File.join(config_src_dir, "postgresql.conf")
  pgsql_dir = File.join(File.dirname(__FILE__), "..", "pgsql")
  setup_sql_path = File.join(pgsql_dir, "setup.sql")

  puts "Setting up '#{config[:data_dir]}'..."
  data_dir = File.expand_path(config[:data_dir])

  # log & socket dirs
  log_dir = File.join(data_dir, "log")
  run_dir = File.join(data_dir, "run")
  FileUtils.mkdir_p log_dir
  FileUtils.mkdir_p run_dir

  pgpool_dir = File.join(data_dir, "pgpool")
  postgres_dir = File.join(data_dir, "postgres")

  @config = setup_params  # used by erb files
  @config[:pgpool_pid_file] = File.join(run_dir, "pgpool.pid")
  @config[:pgpool_status_dir] = run_dir
  @config[:unix_socket_directory] = run_dir

  # pgpool
  FileUtils.mkdir_p pgpool_dir
  %w[pcp.conf.sample pgpool.conf pool_hba.conf pool_passwd].each do |fname|
    # run ERB for each config files
    dst_path = File.join(pgpool_dir, fname)
    src_path = File.join(config_src_dir, fname)
    File.open(dst_path, "w") {|f|
      f.write ERB.new(File.read(src_path)).result
    }
  end

  # postgresql
  # 1. initdb postgres_dir
  setup_cmd "initdb -U pg --no-locale -E UNICODE #{se(postgres_dir)}"
  File.open(File.join(data_dir, "postgres", "postgresql.conf"), "a") {|f|
    f.write ERB.new(File.read(postgresql_conf_path)).result
  }

  # 2. postgres -D postgres_dir -c listen_addresses=
  # start postgres without listening on TCP
  postgres_cmd = "postgres -D #{se(postgres_dir)} -c listen_addresses="
  puts "setup> #{postgres_cmd}"
  pid = spawn postgres_cmd

  # 3. psql < pgsql/setup.sql
  begin
    # waits for spinup of postgres
    sleep 2
    puts "initializing database..."
    sleep 8

    psql_cmd = "psql -h #{se(@config[:unix_socket_directory])} -p #{@config[:backend_port]} -U pg postgres"
    setup_cmd "#{psql_cmd} < #{se(setup_sql_path)}"
  ensure
    # 4. shutdown postgres
    Process.kill(:SIGQUIT, pid)
  end
  Process.waitpid(pid)

  puts <<EOF

************************************************************
**         PostgreSQL protocol gateway for Presto         **
************************************************************

           +------------+   +------------+   +--------+
  client ---> pgpool-II ----> PostgreSQL ----> Presto |
           +------------+   +------------+   +--------+
           rewrite query     proxy           run query

You can run queries on Presto as following:

  * Set presto_server and presto_catalog to pgpool.conf:
    #{config[:data_dir]}/pgpool/pgpool.conf

  * Start pgpool-II:
    $ prestogres -D "#{config[:data_dir]}" pgpool
    (stop> prestogres -D "#{config[:data_dir]}" pgpool stop)

  * Start PostgreSQL:
    $ prestogres -D "#{config[:data_dir]}" pg_ctl start
    (stop> prestogres -D "#{config[:data_dir]}" pg_ctl stop)

  * Connect to pgpool-II:
    $ psql -h localhost -p #{setup_params[:port]} -U pg postgres

  If you run PostgreSQL on Mac OS X, you need to run following commands:

    $ sudo sysctl -w kern.sysv.shmmax=1073741824
    $ sudo sysctl -w kern.sysv.shmall=1073741824

You need to update 'presto_server' and 'presto_catalog' in
pgpool.conf at least.

EOF

when :passwd
  if args.empty?
    usage nil
  end
  unless config[:data_dir]
    usage "-D option is required"
  end

  require "prestogres_config"
  prefix = Prestogres::CONFIG["prefix"]
  binary = File.join(prefix, "bin", 'pg_md5')

  data_dir = config[:data_dir]
  pgpool_conf_path = File.join(data_dir, "pgpool", "pgpool.conf")

  args = ["-f", pgpool_conf_path, "-p", "-m", "-u"].concat(args)

  puts "#{binary} #{args.join(' ')}"
  system binary, *args

  if $?.success?
    puts "Successfully pdated pgpool/pool_passwd file. Edit pgpool/pool_passwd file directly to remove the entry."
  else
    exit $?.exitstatus
  end

when :show_init_sql
  pgsql_dir = File.join(File.dirname(__FILE__), "..", "pgsql")
  STDOUT.write File.read(File.join(pgsql_dir, "setup.sql"))

else
  unless config[:data_dir]
    usage "-D option is required"
  end
  require "prestogres_config"

  if config[:command] == :pgpool
    prefix = Prestogres::CONFIG["prefix"]
    binary = File.join(prefix, "bin", config[:command].to_s)
  else
    # TODO configurable postgres install prefix (by ENV?)
    binary = config[:command].to_s
  end

  pgsql_dir = File.expand_path File.join(File.dirname(__FILE__), "..", "pgsql")
  ENV['PYTHONPATH'] = [pgsql_dir, ENV['PYTHONPATH']].compact.join(':')

  data_dir = config[:data_dir]
  case config[:command]
  when :pgpool
    args = ["--pcp-file", File.join(data_dir, "pgpool", "pcp.conf")].concat(args)
    args = ["--hba-file", File.join(data_dir, "pgpool", "pool_hba.conf")].concat(args)
    args = ["--config-file", File.join(data_dir, "pgpool", "pgpool.conf")].concat(args)
  when :pg_ctl
    args = ["-D", File.join(data_dir, "postgres")].concat(args)
    args = ["-l", File.join(data_dir, "log", "postgres.log")].concat(args)
  when :postgres
    args = ["-D", File.join(data_dir, "postgres")].concat(args)
  end

  puts "PYTHONPATH=#{ENV['PYTHONPATH']}"
  puts "#{binary} #{args.join(' ')}"

  if Kernel.respond_to?(:exec)
    exec binary, *args
  else
    system binary, *args
    exit $?.exitstatus
  end

end