Sequel.require 'adapters/shared/postgres'

module Sequel
  Postgres::CONVERTED_EXCEPTIONS << NativeException
  
  module JDBC
    # Adapter, Database, and Dataset support for accessing a PostgreSQL
    # database via JDBC.
    module Postgres
      # Methods to add to Database instances that access PostgreSQL via
      # JDBC.
      module DatabaseMethods
        extend Sequel::Database::ResetIdentifierMangling
        include Sequel::Postgres::DatabaseMethods
        
        # Add the primary_keys and primary_key_sequences instance variables,
        # so we can get the correct return values for inserted rows.
        def self.extended(db)
          super
          db.send(:initialize_postgres_adapter)
        end

        # See Sequel::Postgres::Adapter#copy_into
        def copy_into(table, opts=OPTS)
          data = opts[:data]
          data = Array(data) if data.is_a?(String)

          if block_given? && data
            raise Error, "Cannot provide both a :data option and a block to copy_into"
          elsif !block_given? && !data
            raise Error, "Must provide either a :data option or a block to copy_into"
          end

          synchronize(opts) do |conn|
            begin
              copy_manager = org.postgresql.copy.CopyManager.new(conn)
              copier = copy_manager.copy_in(copy_into_sql(table, opts))
              if block_given?
                while buf = yield
                  copier.writeToCopy(buf.to_java_bytes, 0, buf.length)
                end
              else
                data.each { |d| copier.writeToCopy(d.to_java_bytes, 0, d.length) }
              end
            rescue Exception => e
              copier.cancelCopy
              raise
            ensure
              unless e
                begin
                  copier.endCopy
                rescue NativeException => e2
                  raise_error(e2)
                end
              end
            end
          end
        end
        
        # See Sequel::Postgres::Adapter#copy_table
        def copy_table(table, opts=OPTS)
          synchronize(opts[:server]) do |conn|
            copy_manager = org.postgresql.copy.CopyManager.new(conn)
            copier = copy_manager.copy_out(copy_table_sql(table, opts))
            begin
              if block_given?
                while buf = copier.readFromCopy
                  yield(String.from_java_bytes(buf))
                end
                nil
              else
                b = ''
                while buf = copier.readFromCopy
                  b << String.from_java_bytes(buf)
                end
                b
              end
            ensure
              raise DatabaseDisconnectError, "disconnecting as a partial COPY may leave the connection in an unusable state" if buf
            end
          end
        end

        private
        
        # Use setNull for nil arguments as the default behavior of setString
        # with nil doesn't appear to work correctly on PostgreSQL.
        def set_ps_arg_nil(cps, i)
          cps.setNull(i, JavaSQL::Types::NULL)
        end

        # Execute the connection configuration SQL queries on the connection.
        def setup_connection(conn)
          conn = super(conn)
          statement(conn) do |stmt|
            connection_configuration_sqls.each{|sql| log_yield(sql){stmt.execute(sql)}}
          end
          conn
        end
      end
      
      # Dataset subclass used for datasets that connect to PostgreSQL via JDBC.
      class Dataset < JDBC::Dataset
        include Sequel::Postgres::DatasetMethods
        APOS = Dataset::APOS
        
        class ::Sequel::JDBC::Dataset::TYPE_TRANSLATOR
          # Convert Java::OrgPostgresqlUtil::PGobject to ruby strings
          def pg_object(v)
            v.to_string
          end
        end

        # Handle conversions of PostgreSQL array instances
        class PGArrayConverter
          # Set the method that will return the correct conversion
          # proc for elements of this array.
          def initialize(meth)
            @conversion_proc_method = meth
            @conversion_proc = nil
          end
          
          # Convert Java::OrgPostgresqlJdbc4::Jdbc4Array to ruby arrays
          def call(v)
            _pg_array(v.array)
          end

          private

          # Handle multi-dimensional Java arrays by recursively mapping them
          # to ruby arrays of ruby values.
          def _pg_array(v)
            v.to_ary.map do |i|
              if i.respond_to?(:to_ary)
                _pg_array(i)
              elsif i
                if @conversion_proc.nil?
                  @conversion_proc = @conversion_proc_method.call(i)
                end
                if @conversion_proc
                  @conversion_proc.call(i)
                else
                  i
                end
              else
                i
              end
            end
          end
        end

        PG_OBJECT_METHOD = TYPE_TRANSLATOR_INSTANCE.method(:pg_object)
      
        # Add the shared PostgreSQL prepared statement methods
        def prepare(type, name=nil, *values)
          ps = to_prepared_statement(type, values)
          ps.extend(JDBC::Dataset::PreparedStatementMethods)
          ps.extend(::Sequel::Postgres::DatasetMethods::PreparedStatementMethods)
          if name
            ps.prepared_statement_name = name
            db.set_prepared_statement(name, ps)
          end
          ps
        end

        private
        
        # Handle PostgreSQL array and object types. Object types are just
        # turned into strings, similarly to how the native adapter treats
        # the types.
        def convert_type_proc(v, ctn=nil)
          case v
          when Java::OrgPostgresqlJdbc4::Jdbc4Array
            if pr = db.conversion_procs[ctn]
              lambda{|x| pr.call(PG_OBJECT_METHOD.call(x))}
            else
              PGArrayConverter.new(method(:convert_type_proc))
            end
          when Java::OrgPostgresqlUtil::PGobject
            if pr = db.conversion_procs[ctn]
              lambda{|x| pr.call(PG_OBJECT_METHOD.call(x))}
            else
              PG_OBJECT_METHOD
            end
          when String
            if pr = db.conversion_procs[ctn]
              pr
            else
              false
            end
          when JAVA_HASH_MAP
            if Sequel.respond_to?(:hstore)
              lambda{|x| Sequel.hstore(HASH_MAP_METHOD.call(x))}
            else
              HASH_MAP_METHOD
            end
          else
            super
          end
        end
        
        # The jdbc/postgresql adapter uses column type oids when determining
        # conversion procs.
        def convert_type_proc_uses_column_info?
          true
        end

        # Use the column type oid as the database specific column info value.
        def convert_type_proc_column_info(meta, i)
          meta.field(i).oid
        end
      
        # Literalize strings similar to the native postgres adapter
        def literal_string_append(sql, v)
          sql << APOS << db.synchronize(@opts[:server]){|c| c.escape_string(v)} << APOS
        end
      end
    end
  end
end