require 'spatial_adapter'
require 'active_record/connection_adapters/postgresql_adapter'

module ActiveRecord::ConnectionAdapters
  class PostgreSQLAdapter
    def postgis_version
      begin
        select_value("SELECT postgis_full_version()").scan(/POSTGIS="([\d\.]*)"/)[0][0]
      rescue ActiveRecord::StatementInvalid
        nil
      end
    end

    def postgis_major_version
      version = postgis_version
      version ? version.scan(/^(\d)\.\d\.\d$/)[0][0].to_i : nil
    end

    def postgis_minor_version
      version = postgis_version
      version ? version.scan(/^\d\.(\d)\.\d$/)[0][0].to_i : nil
    end

    def spatial?
      !postgis_version.nil?
    end

    def supports_geographic?
      postgis_major_version > 1 || (postgis_major_version == 1 && postgis_minor_version >= 5)
    end

    alias :original_native_database_types :native_database_types
    def native_database_types
      original_native_database_types.merge!(SpatialAdapter.geometry_data_types)
    end

    alias :original_quote :quote
    #Redefines the quote method to add behaviour for when a Geometry is encountered
    def quote(value, column = nil)
      if value.kind_of?(GeoRuby::SimpleFeatures::Geometry)
        "'#{value.as_hex_ewkb}'"
      else
        original_quote(value,column)
      end
    end

    def columns(table_name, name = nil) #:nodoc:
      raw_geom_infos = column_spatial_info(table_name)

      column_definitions(table_name).collect do |name, type, default, notnull|
        case type
        when /geography/i
          ActiveRecord::ConnectionAdapters::SpatialPostgreSQLColumn.create_from_geography(name, default, type, notnull == 'f')
        when /geometry/i
          raw_geom_info = raw_geom_infos[name]
          if raw_geom_info.nil?
            # This column isn't in the geometry_columns table, so we don't know anything else about it
            ActiveRecord::ConnectionAdapters::SpatialPostgreSQLColumn.create_simplified(name, default, notnull == "f")
          else
            ActiveRecord::ConnectionAdapters::SpatialPostgreSQLColumn.new(name, default, raw_geom_info.type, notnull == "f", raw_geom_info.srid, raw_geom_info.with_z, raw_geom_info.with_m)
          end
        else
          ActiveRecord::ConnectionAdapters::PostgreSQLColumn.new(name, default, type, notnull == "f")
        end
      end
    end

    def create_table(table_name, options = {})
      # Using the subclassed table definition
      table_definition = ActiveRecord::ConnectionAdapters::PostgreSQLTableDefinition.new(self)
      table_definition.primary_key(options[:primary_key] || ActiveRecord::Base.get_primary_key(table_name.to_s.singularize)) unless options[:id] == false

      yield table_definition if block_given?

      if options[:force] && table_exists?(table_name)
        drop_table(table_name, options)
      end

      create_sql = "CREATE#{' TEMPORARY' if options[:temporary]} TABLE "
      create_sql << "#{quote_table_name(table_name)} ("
      create_sql << table_definition.to_sql
      create_sql << ") #{options[:options]}"

      # This is the additional portion for PostGIS
      unless table_definition.geom_columns.nil?
        table_definition.geom_columns.each do |geom_column|
          geom_column.table_name = table_name
          create_sql << "; " + geom_column.to_sql
        end
      end

      execute create_sql
    end

    alias :original_remove_column :remove_column
    def remove_column(table_name, *column_names)
      column_names = column_names.flatten
      columns(table_name).each do |col|
        if column_names.include?(col.name.to_sym)
          # Geometry columns have to be removed using DropGeometryColumn
          if col.is_a?(::SpatialAdapter::SpatialColumn) && col.spatial? && !col.geographic?
            execute "SELECT DropGeometryColumn('#{table_name}','#{col.name}')"
          else
            original_remove_column(table_name, col.name)
          end
        end
      end
    end

    alias :original_add_column :add_column
    def add_column(table_name, column_name, type, options = {})
      unless SpatialAdapter.geometry_data_types[type].nil?
        geom_column = ActiveRecord::ConnectionAdapters::PostgreSQLColumnDefinition.new(self, column_name, type, nil, nil, options[:null], options[:srid] || -1 , options[:with_z] || false , options[:with_m] || false, options[:geographic] || false)
        if geom_column.geographic
          default = options[:default]
          notnull = options[:null] == false

          execute("ALTER TABLE #{quote_table_name(table_name)} ADD COLUMN #{geom_column.to_sql}")

          change_column_default(table_name, column_name, default) if options_include_default?(options)
          change_column_null(table_name, column_name, false, default) if notnull
        else
          geom_column.table_name = table_name
          execute geom_column.to_sql
        end
      else
        original_add_column(table_name, column_name, type, options)
      end
    end

    # Adds an index to a column.
    def add_index(table_name, column_name, options = {})
      column_names = Array(column_name)
      index_name   = index_name(table_name, :column => column_names)

      if Hash === options # legacy support, since this param was a string
        index_type = options[:unique] ? "UNIQUE" : ""
        index_name = options[:name] || index_name
        index_method = options[:spatial] ? 'USING GIST' : ""
      else
        index_type = options
      end
      quoted_column_names = column_names.map { |e| quote_column_name(e) }.join(", ")
      execute "CREATE #{index_type} INDEX #{quote_column_name(index_name)} ON #{quote_table_name(table_name)} #{index_method} (#{quoted_column_names})"
    end

    # Returns the list of all indexes for a table.
    #
    # This is a full replacement for the ActiveRecord method and as a result
    # has a higher probability of breaking in future releases.
    def indexes(table_name, name = nil)
       schemas = schema_search_path.split(/,/).map { |p| quote(p) }.join(',')

       # Changed from upstread: link to pg_am to grab the index type (e.g. "gist")
       result = query(<<-SQL, name)
         SELECT distinct i.relname, d.indisunique, d.indkey, t.oid, am.amname
           FROM pg_class t, pg_class i, pg_index d, pg_attribute a, pg_am am
         WHERE i.relkind = 'i'
           AND d.indexrelid = i.oid
           AND d.indisprimary = 'f'
           AND t.oid = d.indrelid
           AND t.relname = '#{table_name}'
           AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname IN (#{schemas}) )
           AND i.relam = am.oid
           AND a.attrelid = t.oid
        ORDER BY i.relname
      SQL


      indexes = []

      indexes = result.map do |row|
        index_name = row[0]
        unique = row[1] == 't'
        indkey = row[2].split(" ")
        oid = row[3]
        indtype = row[4]

        # Changed from upstream: need to get the column types to test for spatial indexes
        columns = query(<<-SQL, "Columns for index #{row[0]} on #{table_name}").inject({}) {|attlist, r| attlist[r[1]] = [r[0], r[2]]; attlist}
        SELECT a.attname, a.attnum, t.typname
        FROM pg_attribute a, pg_type t
        WHERE a.attrelid = #{oid}
        AND a.attnum IN (#{indkey.join(",")})
        AND a.atttypid = t.oid
        SQL

        # Only GiST indexes on spatial columns denote a spatial index
        spatial = indtype == 'gist' && columns.size == 1 && (columns.values.first[1] == 'geometry' || columns.values.first[1] == 'geography')

        column_names = indkey.map {|attnum| columns[attnum] ? columns[attnum][0] : nil }
        ActiveRecord::ConnectionAdapters::IndexDefinition.new(table_name, index_name, unique, column_names, spatial)
      end

      indexes
    end

    def disable_referential_integrity(&block) #:nodoc:
      if supports_disable_referential_integrity?() then
        execute(tables_without_postgis.collect { |name| "ALTER TABLE #{quote_table_name(name)} DISABLE TRIGGER ALL" }.join(";"))
      end
      yield
    ensure
      if supports_disable_referential_integrity?() then
        execute(tables_without_postgis.collect { |name| "ALTER TABLE #{quote_table_name(name)} ENABLE TRIGGER ALL" }.join(";"))
      end
    end

    private

    def tables_without_postgis
      tables - %w{ geometry_columns spatial_ref_sys }
    end

    def column_spatial_info(table_name)
      constr = query("SELECT * FROM geometry_columns WHERE f_table_name = '#{table_name}'")

      raw_geom_infos = {}
      constr.each do |constr_def_a|
        raw_geom_infos[constr_def_a[3]] ||= SpatialAdapter::RawGeomInfo.new
        raw_geom_infos[constr_def_a[3]].type = constr_def_a[6]
        raw_geom_infos[constr_def_a[3]].dimension = constr_def_a[4].to_i
        raw_geom_infos[constr_def_a[3]].srid = constr_def_a[5].to_i

        if raw_geom_infos[constr_def_a[3]].type[-1] == ?M
          raw_geom_infos[constr_def_a[3]].with_m = true
          raw_geom_infos[constr_def_a[3]].type.chop!
        else
          raw_geom_infos[constr_def_a[3]].with_m = false
        end
      end

      raw_geom_infos.each_value do |raw_geom_info|
        #check the presence of z and m
        raw_geom_info.convert!
      end

      raw_geom_infos
    end
  end

  class PostgreSQLTableDefinition < TableDefinition
    attr_reader :geom_columns

    def column(name, type, options = {})
      unless (SpatialAdapter.geometry_data_types[type.to_sym].nil? or
              (options[:create_using_addgeometrycolumn] == false))

        column = self[name] || PostgreSQLColumnDefinition.new(@base, name, type)
        column.null = options[:null]
        column.srid = options[:srid] || -1
        column.with_z = options[:with_z] || false
        column.with_m = options[:with_m] || false
        column.geographic = options[:geographic] || false

        if column.geographic
          @columns << column unless @columns.include? column
        else
          # Hold this column for later
          @geom_columns ||= []
          @geom_columns << column
        end
        self
      else
        super(name, type, options)
      end
    end
  end

  class PostgreSQLColumnDefinition < ColumnDefinition
    attr_accessor :table_name
    attr_accessor :srid, :with_z, :with_m, :geographic
    attr_reader :spatial

    def initialize(base = nil, name = nil, type=nil, limit=nil, default=nil, null=nil, srid=-1, with_z=false, with_m=false, geographic=false)
      super(base, name, type, limit, default, null)
      @table_name = nil
      @spatial = true
      @srid = srid
      @with_z = with_z
      @with_m = with_m
      @geographic = geographic
    end

    def sql_type
      if geographic
        type_sql = SpatialAdapter.geometry_data_types[type.to_sym][:name]
        type_sql += "Z" if with_z
        type_sql += "M" if with_m
        # SRID is not yet supported (defaults to 4326)
        #type_sql += ", #{srid}" if (srid && srid != -1)
        type_sql = "geography(#{type_sql})"
        type_sql
      else
        super
      end
    end

    def to_sql
      if spatial && !geographic
        type_sql = SpatialAdapter.geometry_data_types[type.to_sym][:name]
        type_sql += "M" if with_m and !with_z
        if with_m and with_z
          dimension = 4
        elsif with_m or with_z
          dimension = 3
        else
          dimension = 2
        end

        column_sql = "SELECT AddGeometryColumn('#{table_name}','#{name}',#{srid},'#{type_sql}',#{dimension})"
        column_sql += ";ALTER TABLE #{table_name} ALTER #{name} SET NOT NULL" if null == false
        column_sql
      else
        super
      end
    end
  end

  class SpatialPostgreSQLColumn < PostgreSQLColumn
    include SpatialAdapter::SpatialColumn

    def initialize(name, default, sql_type = nil, null = true, srid=-1, with_z=false, with_m=false, geographic = false)
      super(name, default, sql_type, null, srid, with_z, with_m)
      @geographic = geographic
    end

    def geographic?
      @geographic
    end

    #Transforms a string to a geometry. PostGIS returns a HewEWKB string.
    def self.string_to_geometry(string)
      return string unless string.is_a?(String)
      GeoRuby::SimpleFeatures::Geometry.from_hex_ewkb(string) rescue nil
    end

    def self.create_simplified(name, default, null = true)
      new(name, default, "geometry", null)
    end

    def self.create_from_geography(name, default, sql_type, null = true)
      params = extract_geography_params(sql_type)
      new(name, default, sql_type, null, params[:srid], params[:with_z], params[:with_m], true)
    end

    private

    # Add detection of PostGIS-specific geography columns
    def geometry_simplified_type(sql_type)
      case sql_type
      when /geography\(point/i then :point
      when /geography\(linestring/i then :line_string
      when /geography\(polygon/i then :polygon
      when /geography\(multipoint/i then :multi_point
      when /geography\(multilinestring/i then :multi_line_string
      when /geography\(multipolygon/i then :multi_polygon
      when /geography\(geometrycollection/i then :geometry_collection
      when /geography/i then :geometry
      else
        super
      end
    end

    def self.extract_geography_params(sql_type)
      params = {
        :srid => 0,
        :with_z => false,
        :with_m => false
      }
      if sql_type =~ /geography(?:\((?:\w+?)(Z)?(M)?(?:,(\d+))?\))?/i
        params[:with_z] = $1 == 'Z'
        params[:with_m] = $2 == 'M'
        params[:srid]   = $3.to_i
      end
      params
    end
  end
end
