require("date")
require("ostruct")
require("optparse")
require("pathname")
require("etc")
require("cipherstash-pg")
class PGWarehouse
  def initialize(opts)
    @opts = opts
    @db = CipherStashPG.connect(:dbname => opts.database, :host => opts.host, :port => opts.port, :user => opts.user, :password => opts.pass, :sslmode => "prefer")
    @db.exec(("SET search_path TO %s" % [opts.schema])) if opts.schema
    @relations = self.relations
  end
  
  attr_reader(:db)
  
  public
  
  def migrate
    if @relations.empty? then
      $stderr.puts("No tables were found for warehousing.")
      return
    end
    $stderr.puts(("Found %d relation%s to move." % [relations.length, (relations.length == 1) ? ("") : ("s")]))
    @relations.sort_by { |_, v| v[:name] }.each do |_, val|
      $stderr.print(("  - Moving table '%s' to '%s'... " % [val[:name], @opts.tablespace]))
      if @opts.dryrun then
        $stderr.puts("(not really)")
      else
        age = self.timer do
          db.exec(("ALTER TABLE %s SET TABLESPACE %s;" % [val[:name], @opts.tablespace]))
        end
        puts(age)
      end
      val[:indexes].each do |idx|
        $stderr.print(("      - Moving index '%s' to '%s'... " % [idx, @opts.tablespace]))
        if @opts.dryrun then
          $stderr.puts("(not really)")
        else
          age = self.timer do
            db.exec(("ALTER INDEX %s SET TABLESPACE %s;" % [idx, @opts.tablespace]))
          end
          puts(age)
        end
      end
    end
  end
  
  protected
  
  def relations
    return @relations if @relations
    relations = {}
    query = "\n\t\t\tSELECT c.oid AS oid,\n\t\t\t\tc.relname AS name,\n\t\t\t\tc.relkind AS kind,\n\t\t\t\tt.spcname AS tspace\n\t\t\tFROM pg_class AS c\n\t\t\tLEFT JOIN pg_namespace n ON n.oid = c.relnamespace\n\t\t\tLEFT JOIN pg_tablespace t ON t.oid = c.reltablespace\n\t\t\tWHERE c.relkind = 'r' "
    (query << "AND n.nspname='#{@opts.schema}'") if @opts.schema
    self.db.exec(query) do |res|
      res.each do |row|
        relations[row["oid"]] = { :name => row["name"], :tablespace => row["tspace"], :indexes => ([]), :parent => nil }
      end
    end
    db.exec("SELECT inhrelid AS oid, inhparent AS parent FROM pg_inherits") do |res|
      res.each { |row| relations[row["oid"]][:parent] = row["parent"] }
    end
    relations.reject! do |oid, val|
      begin
        (val[:parent].nil? or ((val[:tablespace] == @opts.tablespace) or ((val[:name] == Time.now.strftime(@opts.format)) or (not DateTime.strptime(val[:name], @opts.format)))))
      rescue ArgumentError
        true
      end
    end
    query = "\n\t\t\tSELECT c.oid AS oid,\n\t\t\t\ti.indexname AS name\n\t\t\tFROM pg_class AS c\n\t\t\tINNER JOIN pg_indexes AS i\n\t\t\t\tON i.tablename = c.relname "
    (query << "AND i.schemaname='#{@opts.schema}'") if @opts.schema
    db.exec(query) do |res|
      res.each do |row|
        (relations[row["oid"]][:indexes] << row["name"]) if relations[row["oid"]]
      end
    end
    return relations
  end
  
  def timer
    start = Time.now
    yield
    age = (Time.now - start)
    diff = age
    secs = (diff % 60)
    diff = ((diff - secs) / 60)
    mins = (diff % 60)
    diff = ((diff - mins) / 60)
    hour = (diff % 24)
    return ("%02d:%02d:%02d" % [hour, mins, secs])
  end
end
def parse_args(args)
  options = OpenStruct.new
  options.database = Etc.getpwuid(Process.uid).name
  options.host = "127.0.0.1"
  options.port = 5432
  options.user = Etc.getpwuid(Process.uid).name
  options.sslmode = "prefer"
  options.tablespace = "warehouse"
  opts = OptionParser.new do |opts|
    opts.banner = "Usage: #{$0} [options]"
    opts.separator("")
    opts.separator("Connection options:")
    opts.on("-d", "--database DBNAME", "specify the database to connect to (default: \"#{options.database}\")") do |db|
      options.database = db
    end
    opts.on("-h", "--host HOSTNAME", "database server host") do |host|
      options.host = host
    end
    opts.on("-p", "--port PORT", Integer, "database server port (default: \"#{options.port}\")") do |port|
      options.port = port
    end
    opts.on("-n", "--schema SCHEMA", String, "operate on the named schema only (default: none)") do |schema|
      options.schema = schema
    end
    opts.on("-T", "--tablespace SPACE", String, "move old tables to this tablespace (default: \"#{options.tablespace}\")") do |tb|
      options.tablespace = tb
    end
    opts.on("-F", "--tableformat FORMAT", String, "The naming format (strftime) for the inherited tables (default: none)") do |format|
      options.format = format
    end
    opts.on("-U", "--user NAME", "database user name (default: \"#{options.user}\")") do |user|
      options.user = user
    end
    opts.on("-W", "force password prompt") do |pw|
      print("Password: ")
      begin
        (system("stty -echo")
        options.pass = gets.chomp)
      ensure
        (system("stty echo")
        puts)
      end
    end
    opts.separator("")
    opts.separator("Other options:")
    opts.on_tail("--dry-run", "don't actually do anything") do
      options.dryrun = true
    end
    opts.on_tail("--help", "show this help, then exit") do
      $stderr.puts(opts)
      exit
    end
    opts.on_tail("--version", "output version information, then exit") do
      puts(Stats::VERSION)
      exit
    end
  end
  opts.parse!(args)
  return options
end
if ("(string)" == $0) then
  opts = parse_args(ARGV)
  unless opts.format then
    raise(ArgumentError, "A naming format (-F) is required.")
  end
  $stdout.sync = true
  PGWarehouse.new(opts).migrate
end