require 'bricolage/datasource'
require 'bricolage/s3datasource'
require 'bricolage/sqlstatement'
require 'bricolage/commandutils'
require 'bricolage/postgresconnection'
require 'bricolage/vacuumlock'
require 'bricolage/exception'
require 'bricolage/logger'
require 'pathname'

module Bricolage

  class PSQLDataSource < DataSource
    declare_type 'psql'

    include CommandUtils

    DEFAULT_RETRY_LIMIT = 3

    def initialize(
        host: 'localhost',
        port: 5439,
        database: 'dev',
        username: ENV['LOGNAME'],
        password: nil,
        pgpass: nil,
        encoding: nil,
        psql: 'psql',
        sql_log_level: nil,
        query_sql_log_level: nil,
        update_sql_log_level: nil,
        tmpdir: Dir.tmpdir)
      @host = host
      @port = port
      @database = database
      @user = username
      @password = password
      @pgpass = pgpass
      @encoding = encoding
      @psql = psql
      @sql_log_level = Logger.intern_severity(sql_log_level || Logger::DEBUG)
      @query_sql_log_level = Logger.intern_severity(query_sql_log_level || sql_log_level || Logger::DEBUG)
      @update_sql_log_level = Logger.intern_severity(update_sql_log_level || sql_log_level || Logger::INFO)
      @tmpdir = tmpdir
      @connection_pool = []
      raise ParameterError, "missing psql host" unless @host
      raise ParameterError, "missing psql port" unless @port
      raise ParameterError, "missing psql database" unless @database
      raise ParameterError, "missing psql username" unless @user
      unless @pgpass or @password
        raise ParameterError, "missing psql password"
      end
    end

    attr_reader :host
    attr_reader :port
    attr_reader :database
    attr_reader :user

    attr_reader :sql_log_level
    attr_reader :query_sql_log_level
    attr_reader :update_sql_log_level

    def new_task
      PSQLTask.new(self)
    end

    def open_for_batch
      # do not call #open
      yield
    end

    def execute(source, options = [])
      make_tmpfile(source, tmpdir: @tmpdir) {|path|
        st = command @psql, "--no-psqlrc", "--host=#{@host}", "--port=#{@port}",
            "--username=#{@user}", @database,
            '--echo-all',
            '-v', 'ON_ERROR_STOP=true',
            '-f', path,
            '--no-password',
            *options,
            env: get_psql_env
        unless st.success?
          begin
            msg = LogLocator.slice_last_stderr(/^psql:.*?:\d+: ERROR: (.*)/, 1)
          rescue IOError => ex
            # slice_last_stderr may fail if stderr is not a file
            logger.error ex.message
            msg = nil
          end
        end
        JobResult.for_process_status(st, msg)
      }
    end

    def get_psql_env
      env = {}
      if @pgpass
        env["PGPASSFILE"] = @pgpass
      elsif @password
        env["PGPASSWORD"] = @password
      end
      env
    end

    #
    # Ruby Library Interface
    #

    def password
      # FIXME: same user must not exist
      @password ||= read_password_from_pgpass(@pgpass, @user)
    end

    def read_password_from_pgpass(path, user)
      File.read(path).slice(/:#{user}:([^:\r\n]+)$/, 1) or
          raise ParameterError, "could not read password: #{path}, #{user}"
    end

    def open
      retries = (ENV['BRICOLAGE_OPEN_RETRY_LIMIT'] || DEFAULT_RETRY_LIMIT).to_i
      begin
        conn = PostgresConnection.open_data_source(self)
        conn.execute_query('select 1') {}
      rescue PG::ConnectionBad, PG::UnableToSend => ex
        conn.close rescue nil
        retries -= 1
        if retries >= 0
          logger.warn "Could not open postgres connection; retry: #{ex.message}"
          sleep 1
          retry
        else
          raise
        end
      end
      if block_given?
        begin
          yield conn
        ensure
          conn.close
        end
      else
        return conn
      end
    end

    def open_shared_connection
      raise ParameterError, 'open_shared_connection require block' unless block_given?
      conn = nil
      until conn
        if conn_tmp = @connection_pool.shift
          begin
            conn_tmp.query('select 1') {}
          rescue PG::ConnectionBad, PG::UnableToSend
            # retry
          else
            # no exception occured
            conn = conn_tmp
            conn_tmp = nil
          ensure
            if conn_tmp
              conn_tmp.close
              conn_tmp = nil
            end
          end
        else
          # Get a fresh connection instead of pooled connections.
          conn = open
        end
      end
      begin
        yield conn
      ensure
        @connection_pool.push(conn)
      end
    end

    # not MT-safe
    def clear_connection_pool
      @connection_pool.map(&:close)
      @connection_pool = []
    end

    def query_batch(query, batch_size = 5000, &block)
      open {|conn| conn.query_batch(query, batch_size, &block) }
    end

    def drop_table(name)
      open {|conn| conn.drop_table(name) }
    end

    def drop_table_force(name)
      open {|conn| conn.drop_table_force(name) }
    end

    def select(table, &block)
      open {|conn| conn.select(table, &block) }
    end

    include VacuumLock

    def vacuum(table)
      serialize_vacuum {
        open {|conn| conn.vacuum(table) }
      }
    end

    def vacuum_sort_only(table)
      serialize_vacuum {
        open {|conn| conn.vacuum_sort_only(table) }
      }
    end

    def analyze(table)
      open {|conn| conn.analyze(table) }
    end
  end

  class DataSource   # reopen
    def redshift_loader_source?
      false
    end
  end

  class S3DataSource   # reopen
    def redshift_loader_source?
      true
    end
  end

  # We don't support dynamodb still now

  class PSQLTask < DataSourceTask
    def exec(stmt)
      add SQLAction.new(stmt)
    end

    def each_statement
      each_action do |action|
        yield action.statement
      end
    end

    class SQLAction < Action
      def initialize(stmt)
        @statement = stmt
      end

      attr_reader :statement

      def_delegator '@statement', :bind
    end

    # override
    def source
      buf = StringIO.new
      buf.puts '\timing on'
      each_statement do |stmt|
        buf.puts
        buf.puts "/* #{stmt.location} */" if stmt.location
        buf.puts stmt.stripped_source
      end
      buf.string
    end

    # override
    def run
      VacuumLock.using {
        @ds.execute source
      }
    end

    def run_explain
      @ds.execute explain_source
    end

    def explain_source
      buf = StringIO.new
      each_statement do |stmt|
        buf.puts
        buf.puts "-- #{stmt.location}" if stmt.location
        if support_explain?(stmt.kind)
          buf.puts "explain #{stmt.stripped_source}"
        else
          buf.puts "/* #{stmt.stripped_source} */"
        end
      end
      buf.string
    end

    def support_explain?(statement_kind)
      case statement_kind
      when 'select', 'insert', 'update', 'delete' then true
      else false
      end
    end

    def create_dummy_table(target)
      exec SQLStatement.for_string(
        "create table if not exists #{target} (x int);\n"
      )
    end

    def drop(target_table)
      exec SQLStatement.for_string("drop table #{target_table} cascade;")
    end

    def drop_if(enabled)
      drop '${dest_table}' if enabled
    end

    def drop_obj_force(type, name)
      exec SQLStatement.for_string(
        "drop #{type} if exists #{name} cascade;\n"
      )
    end

    def drop_force(target_table)
      drop_obj_force('table', target_table)
    end

    def drop_view_force(target_view)
      drop_obj_force('view', target_view)
    end

    def drop_force_if(enabled)
      drop_force('${dest_table}') if enabled
    end

    def drop_view_force_if(enabled)
      drop_view_force('${dest_table}') if enabled
    end

    def truncate_if(enabled, target = '${dest_table}')
      exec SQLStatement.for_string("truncate #{target};") if enabled
    end

    def rename_table(src, dest)
      exec SQLStatement.for_string("alter table #{src} rename to #{dest};")
    end

    def vacuum_if(enable_vacuum, enable_vacuum_sort, target = '${dest_table}')
      if enable_vacuum
        serialize_vacuum {
          exec SQLStatement.for_string("vacuum #{target};")
        }
      elsif enable_vacuum_sort
        serialize_vacuum {
          exec SQLStatement.for_string("vacuum sort only #{target};")
        }
      end
    end

    include VacuumLock

    def serialize_vacuum   # override
      exec SQLStatement.for_string psql_serialize_vacuum_begin
      yield
      exec SQLStatement.for_string psql_serialize_vacuum_end
    end

    def analyze_if(enabled, target = '${dest_table}')
      exec SQLStatement.for_string("analyze #{target};") if enabled
    end

    def grant(privilege:, on:, to:)
      exec SQLStatement.for_string("grant #{privilege} on #{on} to #{to};")
    end

    GRANT_OPTS = %w[privilege to]

    def grant_if(opts, target)
      return unless opts
      return if opts.empty?
      unknown_keys = opts.keys - GRANT_OPTS
      raise ParameterError, "unknown grant options: #{unknown_keys.inspect}" unless unknown_keys.empty?
      missing_keys = GRANT_OPTS - opts.keys
      raise ParameterError, %Q(missing grant options: #{missing_keys.inspect}) unless missing_keys.empty?
      args = {on: target}
      opts.each do |k, v|
        args[k.intern] = v
      end
      grant(**args)
    end

    def transaction
      exec SQLStatement.for_string('begin transaction;')
      yield
      exec SQLStatement.for_string('commit;')
    end

    def load(src_ds, src_path, dest_table, format, jsonpath, opts)
      exec SQLStatement.for_string(copy_statement(src_ds, src_path, dest_table, format, jsonpath, opts))
    end

    def copy_statement(src_ds, src_path, dest_table, format, jsonpath, opts)
      unless src_ds.redshift_loader_source?
        raise ParameterError, "input data source does not support redshift as bulk loading source: #{src_ds.name}"
      end
      buf = StringIO.new
      buf.puts "copy #{dest_table}"
      buf.puts "from '#{src_ds.url(src_path)}'"
      buf.puts "credentials '#{src_ds.credential_string}'"
      buf.puts "region '#{src_ds.region}'"
      buf.puts format_option(format, src_ds, jsonpath)
      opts.each do |opt|
        buf.puts opt.to_s
      end
      buf.puts ';'
      buf.string
    end

    def format_option(fmt, src_ds, jsonpath)
      case fmt
      when 'tsv'
        %q(delimiter '\t')
      when 'csv'
        %q(delimiter ',')
      when 'json'
        "json '#{json_param(jsonpath)}'"
      else
        raise ParameterError, "unsupported format: #{fmt}"
      end
    end

    def json_param(jsonpath)
      case jsonpath
      when nil
        'auto'
      when %r{\As3://}
        jsonpath
      else
        src_ds.url(jsonpath)
      end
    end

    def unload(stmt, dest_ds, dest_path, format, opts)
      exec unload_statement(stmt, dest_ds, dest_path, format, opts)
    end

    def unload_statement(stmt, dest_ds, dest_path, format, opts)
      buf = StringIO.new
      buf.puts "unload ('#{format_query(stmt.stripped_raw_content)}')"
      buf.puts "to '#{dest_ds.url(dest_path)}'"
      buf.puts "credentials '#{dest_ds.credential_string}'"
      buf.puts "region '#{dest_ds.region}'"
      buf.puts unload_format_option(format, dest_ds)
      opts.each do |opt|
        buf.puts opt.to_s
      end
      buf.puts ';'
      res = StringResource.new(buf.string, stmt.location)
      SQLStatement.new(res, stmt.declarations)
    end

    def unload_format_option(format, ds)
      case format
      when 'tsv'
        %q(delimiter '\t')
      when 'csv'
        %q(delimiter ',')
      when 'parquet'
        'format as parquet'
      when 'json'
        'format as json'
      else
        raise ParameterError, "unsupported format: #{fmt}"
      end
    end

    def format_query(query)
      query.gsub(/^--.*/, '').strip.gsub(/[ \t]*\n[ \t]*/, ' ').gsub(/\\/,"\\\\\\\\").gsub("'", "\\\\'")
    end
  end

  class PSQLLoadOptions
    class << self
      def parse(opts)
        case opts
        when Hash then filter_values(opts)
        when String then parse_string(opts)   # FIXME: remove
        else raise ParameterError, "unsupported value type for load options: #{opts.class}"
        end
      end

      private

      def filter_values(opts)
        list = []
        opts.each do |key, value|
          case value
          when String, Integer, true, false, nil
            list.push Option.new(key.to_s, value)
          else
            raise ParameterError, ""
          end
        end
        new(list)
      end

      def parse_string(str)
        return new unless str
        list = []
        str.split(',').each do |opt_pair|
          opt, value = opt_pair.strip.split('=', 2)
          list.push Option.new(opt, parse_value(value, opt))
        end
        new(list)
      end

      def parse_value(value, opt)
        case value
        when nil, 'true', true then true
        when 'false', false then false
        when /\A\d+\z/ then value.to_i
        when String then value
        else
          raise ParameterError, "unsupported load option value for #{opt}: #{value.inspect}"
        end
      end
    end

    def initialize(opts = [])
      @opts = opts
    end

    def key?(name)
      n = name.to_s
      @opts.any? {|opt| opt.name == n }
    end

    def each(&block)
      @opts.each(&block)
    end

    def []=(name, value)
      n = name.to_s
      delete n
      @opts.push Option.new(n, value)
    end

    def delete(name)
      @opts.reject! {|opt| opt.name == name }
    end

    def merge(pairs)
      h = {}
      @opts.each do |opt|
        h[opt.name] = opt
      end
      pairs.each do |key, value|
        h[key] = Option.new(key, value)
      end
      self.class.new(h.values)
    end

    def to_s
      buf = StringIO.new
      each do |opt|
        buf.puts opt
      end
      buf.string
    end

    def provide_defaults(src_ds)
    end

    class Option
      def initialize(name, value)
        @name = name
        @value = value
      end

      attr_reader :name
      attr_reader :value

      # Make polymorphic?
      def to_s
        if @name == 'encoding'   # encoding utf16le
          "#{@name} #{@value}"
        else
          case @value
          when true      # acceptanydate
            @name
          when false     # compupdate false
            "#{@name} false"
          when 'on', 'off'
            "#{@name} #{@value}"
          when String    # json 'auto'
            "#{@name} '#{@value}'"
          when Integer   # maxerror 10
            "#{@name} #{@value}"
          when nil       # (explicitly disable boolean options)
            ''
          else
            raise ParameterError, "unsupported option value type for #{@name}: #{@value.inspect}"
          end
        end
      end
    end
  end

end
