# 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 "fileutils"
require "socket"
require "tempfile"
require "time"

require "arrow-flight-sql"

module Helper
  module CommandRunnable
    def spawn_process(*args)
      env = {
        "LC_ALL" => "C",
        "PGCLIENTENCODING" => "UTF-8",
      }
      if args.first.is_a?(Hash)
        env.merge!(args.shift)
      end
      output_read, output_write = IO.pipe
      error_read, error_write = IO.pipe
      options = {
        :out => output_write,
        :err => error_write,
      }
      pid = spawn(env, *args, options)
      output_write.close
      error_write.close
      [pid, output_read, error_read]
    end

    def read_command_output(input)
      return "" unless IO.select([input], nil, nil, 0)
      begin
        data = input.readpartial(4096).gsub(/\r\n/, "\n")
        data.force_encoding("UTF-8")
        data
      rescue EOFError
        ""
      end
    end

    def run_command(*args)
      pid, output_read, error_read = spawn_process(*args)
      output = +""
      error = +""
      status = nil
      timeout = 1
      loop do
        readables, = IO.select([output_read, error_read], nil, nil, timeout)
        if readables
          timeout = 0
          readables.each do |readable|
            if readable == output_read
              output << read_command_output(output_read)
            else
              error << read_command_output(error_read)
            end
          end
        else
          timeout = 1
        end
        _, status = Process.waitpid2(pid, Process::WNOHANG)
        break if status
      end
      output << read_command_output(output_read)
      error << read_command_output(error_read)
      unless status.success?
        command_line = args.join(" ")
        message = "failed to run: #{command_line}\n"
        message << "output:\n"
        message << output
        message << "error:\n"
        message << error
        raise message
      end
      [output, error]
    end
  end

  class PostgreSQL
    include CommandRunnable

    attr_reader :dir
    attr_reader :address
    attr_reader :port
    attr_reader :flight_sql_port
    attr_reader :flight_sql_uri
    attr_reader :user
    attr_reader :password
    def initialize(base_dir)
      @base_dir = base_dir
      @dir = nil
      @log_base_name = "postgresql.log"
      @log_path = nil
      @address = "127.0.0.1"
      @port = nil
      @flight_sql_port = nil
      @flight_sql_uri = nil
      @user = "arrow-flight-sql-test"
      @password = "Passw0rd!"
      @pid = nil
      @running = false
    end

    def running?
      @running
    end

    def initdb(shared_preload_libraries: [],
               max_n_rows_per_record_batch: nil,
               db_path: "db",
               port: 25432,
               flight_sql_port: 35432)
      @dir = File.join(@base_dir, db_path)
      @log_path = File.join(@dir, "log", @log_base_name)
      @port = port
      @pgpass = Tempfile.new("arrow-flight-sql-test-pgpass")
      @pgpass.puts("#{@address}:#{@port}:*:#{@user}:#{@password}")
      @pgpass.close
      @flight_sql_port = flight_sql_port
      if use_tls?
        @flight_sql_uri = "grpc+tls://#{@address}:#{@flight_sql_port}"
      else
        @flight_sql_uri = "grpc://#{@address}:#{@flight_sql_port}"
      end
      Tempfile.create("arrow-flight-sql-test-password") do |password|
        password.print(@password)
        password.close
        run_command("initdb",
                    "--locale", "C",
                    "--encoding", "UTF-8",
                    "--username", @user,
                    "--pwfile", password.path,
                    "-D", @dir)
      end
      prepare_tls if use_tls?
      postgresql_conf = File.join(@dir, "postgresql.conf")
      File.open(postgresql_conf, "a") do |conf|
        conf.puts("listen_addresses = '#{@address}'")
        conf.puts("port = #{@port}")
        conf.puts("unix_socket_directories = ''")
        if use_tls?
          conf.puts("ssl = on")
          conf.puts("ssl_ca_file = 'root.crt'")
        end
        conf.puts("logging_collector = on")
        conf.puts("log_filename = '#{@log_base_name}'")
        conf.puts("log_min_messages = debug5") if ENV["AFS_VERBOSE"] == "yes"
        conf.puts("shared_preload_libraries = " +
                  "'#{shared_preload_libraries.join(",")}'")
        conf.puts("arrow_flight_sql.uri = '#{@flight_sql_uri}'")
        if max_n_rows_per_record_batch
          conf.puts("arrow_flight_sql.max_n_rows_per_record_batch = " +
                    "#{max_n_rows_per_record_batch}")
        end
        yield(conf) if block_given?
      end
      pg_hba_conf = File.join(@dir, "pg_hba.conf")
      pg_hba = File.read(pg_hba_conf)
      pg_hba.gsub!(/^(host.+)trust$/, "\\1password")
      File.write(pg_hba_conf, pg_hba)
    end

    def start
      begin
        run_command("pg_ctl", "start",
                    "-w",
                    "-D", @dir)
      rescue => error
        error.message << "\nPostgreSQL log:\n#{read_log}"
        raise
      end
      loop do
        begin
          TCPSocket.open(@host, @port) do
          end
        rescue SystemCallError
          sleep(0.1)
        else
          break
        end
      end
      @running = true
      pid_path = File.join(@dir, "postmaster.pid")
      if File.exist?(pid_path)
        first_line = File.readlines(pid_path, chomp: true)[0]
        begin
          @pid = Integer(first_line, 10)
        rescue ArgumentError
        end
      end
    end

    def stop
      return unless running?
      begin
        run_command("pg_ctl", "stop",
                    "-D", @dir,
                    "-t", "60")
      rescue
        if @pid
          Process.kill(:KILL, @pid)
          @pid = nil
          @running = false
        end
        raise
      else
        @pid = nil
        @running = false
      end
    end

    def psql(db, sql)
      output, error = run_command({
                                    "PGPASSFILE" => @pgpass.path,
                                  },
                                  "psql",
                                  "--host", @address,
                                  "--port", @port.to_s,
                                  "--username", @user,
                                  "--dbname", db,
                                  "--echo-all",
                                  "--no-password",
                                  "--no-psqlrc",
                                  "--command", sql)
      [output, error]
    end

    def flight_client
      @flight_client ||=
        ArrowFlight::Client.new(@flight_sql_uri, flight_client_options)
    end

    def flight_client_options
      @flight_client_options ||= create_flight_client_options
    end

    def flight_sql_client
      @flight_sql_client ||= ArrowFlightSQL::Client.new(flight_client)
    end

    def read_log
      return "" unless File.exist?(@log_path)
      File.read(@log_path)
    end

    private
    def windows?
      /mingw|mswin|cygwin/.match?(RUBY_PLATFORM)
    end

    def use_tls?
      return false if windows?
      ArrowFlight::ClientOptions.method_defined?(:tls_root_certificates=)
    end

    def create_flight_client_options
      options = ArrowFlight::ClientOptions.new
      if use_tls?
        options.tls_root_certificates = File.read(File.join(@dir, "root.crt"))
        options.override_host_name = "server.example.com"
      end
      options
    end

    def prepare_tls
      prepare_tls_sh = File.join(__dir__, "..", "..", "dev", "prepare-tls.sh")
      prepare_tls_sh = File.expand_path(prepare_tls_sh)
      Dir.chdir(@dir) do
        run_command(prepare_tls_sh,
                    "root.example.com",
                    "server.example.com",
                    "client.example.com")
      end
    end
  end

  module Sandbox
    include CommandRunnable

    class << self
      def included(base)
        base.module_eval do
          setup :setup_tmp_dir

          setup :setup_db

          setup :setup_postgres

          setup :setup_test_db
        end
      end
    end

    def psql(db, sql)
      @postgresql.psql(db, sql)
    end

    def run_sql(sql)
      psql(@test_db_name, sql)
    end

    def flight_client
      @postgresql.flight_client
    end

    def flight_sql_client
      @postgresql.flight_sql_client
    end

    def setup_tmp_dir
      memory_fs = "/dev/shm"
      if File.exist?(memory_fs)
        tmp_dir = memory_fs
      else
        tmp_dir = nil
      end
      Dir.mktmpdir("arrow-flight-sql-", tmp_dir) do |dir|
        @tmp_dir = dir
        begin
          yield
        ensure
          debug_dir = ENV["AFS_TEST_DEBUG_DIR"]
          if debug_dir and File.exist?(@tmp_dir)
            FileUtils.rm_rf(debug_dir)
            FileUtils.mv(@tmp_dir, debug_dir)
          end
        end
      end
    end

    def setup_db
      @postgresql = PostgreSQL.new(@tmp_dir)
      options = {
        shared_preload_libraries: shared_preload_libraries,
        max_n_rows_per_record_batch: max_n_rows_per_record_batch,
      }
      @postgresql.initdb(**options)
      yield
    end

    def shared_preload_libraries
      ["arrow_flight_sql"]
    end

    def max_n_rows_per_record_batch
      nil
    end

    def start_postgres
      @postgresql.start
    end

    def stop_postgres
      @postgresql.stop
    end

    def setup_postgres
      start_postgres
      begin
        yield
      ensure
        if @postgresql
          begin
            stop_postgres
          ensure
            puts(@postgresql.read_log) unless passed?
          end
        end
      end
    end

    def create_db(postgresql, db_name)
      postgresql.psql("postgres", "CREATE DATABASE #{db_name}")
      postgresql.psql(db_name, "CHECKPOINT")
    end

    def setup_test_db
      @test_db_name = "test"
      create_db(@postgresql, @test_db_name)
      result, = run_sql("SELECT oid FROM pg_catalog.pg_database " +
                        "WHERE datname = current_database()")
      oid = result.lines[3].strip
      @test_db_dir = File.join(@postgresql.dir, "base", oid)
      yield
    end
  end
end
