module ReeDao::DatasetExtensions::InstanceMethods

def all(mapper: nil)
  with_mapper(mapper).__original_all
end

def delete(hash_or_entity = nil)
  return __original_delete if hash_or_entity.nil?
  return where(hash_or_entity).__original_delete if hash_or_entity.is_a?(Hash)
  key_condition = prepare_key_condition_from_entity(hash_or_entity)
  where(key_condition).__original_delete
end

def delete_all
  where({}).__original_delete
end

def extract_changes(entity, hash)
  return hash unless entity.instance_variable_defined?(PERSISTENCE_STATE_VARIABLE)
  changes = {}
  persistence_state = entity.instance_variable_get(PERSISTENCE_STATE_VARIABLE)
  hash.each do |column, value|
    previous_column_value = persistence_state[column]
    if persistence_state.has_key?(column)
      if previous_column_value != value || value.respond_to?(:each)
        changes[column] = value
      end
    end
  end
  changes
end

def extract_primary_key(entity)
  if primary_key.is_a?(Array)
    primary_key.map do |key|
      entity.send(key)
    end.join("_")
  elsif entity.instance_variable_defined?(:"@#{primary_key}")
    entity.send(primary_key)
  end
end

def find(id, mapper: nil)
  where(primary_key => id).first(mapper: mapper)
end

def first(mapper: nil)
  with_mapper(mapper).__original_first
end

def ids(ids_list)
  where(id: ids_list)
end

def import_all(entities, batch_size: IMPORT_BATCH_SIZE)
  return if entities.empty?
  mapper = opts[:schema_mapper]
  columns = self.columns
  raw = {}
  columns.delete(:id)
  columns.delete(:row)
  data = entities.map do |entity|
    hash = mapper.db_dump(entity)
    raw[entity] = hash
    columns.map { hash[_1] }
  end
  ids = import(
    columns, data, slice: batch_size, return: :primary_key
  )
  ids.each_with_index do |id, index|
    entity = entities[index]
    raw_data = raw[entity]
    set_entity_primary_key(entity, raw_data, id)
    set_persistence_state(entity, raw_data)
  end
  nil
end

def last(mapper: nil)
  with_mapper(mapper).__original_last
end

def naked_first
  __original_first
end

def naked_last
  __original_last
end

def prepare_key_condition_from_entity(entity)
  key_condition = {}
  if primary_key.is_a?(Array)
    primary_key.each do |key_part|
      key_part_value = entity.send(key_part)
      raise ArgumentError, "entity's primary key can't be nil, got nil for #{key_part}" unless key_part_value
      key_part_value = key_part_value.number if key_part_value.is_a?(ReeEnum::Value)
      key_condition[key_part] = key_part_value
    end
  elsif primary_key.is_a?(Symbol)
    key_value = entity.send(primary_key)
    raise ArgumentError, "entity's primary key can't be nil, got nil for #{primary_key}" unless key_value
    key_condition[primary_key] = key_value
  else
    raise StandardError, "primary key should be array or symbol"
  end
  key_condition
end

def primary_key
  opts[:primary_key] || :id
end

def put(entity)
  raw = opts[:schema_mapper].db_dump(entity)
  remove_null_primary_key(raw)
  key = insert(raw)
  set_entity_primary_key(entity, raw, key)
  set_persistence_state(entity, raw)
  entity
end

def put_with_conflict(entity, conflict_opts = {})
  raw = opts[:schema_mapper].db_dump(entity)
  remove_null_primary_key(raw)
  if conflict_opts.key?(:update) && conflict_opts[:update].empty?
    conflict_opts.delete(:update)
  elsif conflict_opts[:update].is_a?(Array)
    update = {}
    conflict_opts[:update].each do |column|
      update[column] = raw[column]
    end
    conflict_opts[:update] = update
  end
  key = insert_conflict(conflict_opts).insert(raw)
  set_entity_primary_key(entity, raw, key)
  set_persistence_state(entity, raw)
  where(primary_key => key).first
end

def remove_null_primary_key(raw)
  return if primary_key.is_a?(Array)
  if raw.has_key?(primary_key) && raw[primary_key].nil?
    raw.delete(primary_key)
  else
    str_key = primary_key.to_s
    if raw.has_key?(str_key) && raw[str_key].nil?
      raw.delete(str_key)
    end
  end
end

def set_entity_primary_key(entity, raw, key)
  if key && !primary_key.is_a?(Array)
    if entity.respond_to?("#{primary_key}=")
      entity.send("#{primary_key}=", key)
    elsif entity.respond_to?("set_#{primary_key}")
      entity.send("set_#{primary_key}", key)
    else
      entity.instance_variable_set("@#{primary_key}", key)
    end
    raw[primary_key] = key
  end
end

def set_persistence_state(entity, raw)
  if !entity.is_a?(Integer) && !entity.is_a?(Symbol)
    entity.instance_variable_set(PERSISTENCE_STATE_VARIABLE, raw)
  end
end

def table_name
  first_source_table
end

def update(hash_or_entity)
  return __original_update(hash_or_entity) if !opts[:schema_mapper]
  return __original_update(hash_or_entity) if hash_or_entity.is_a?(Hash)
  raw = opts[:schema_mapper].db_dump(hash_or_entity)
  raw = extract_changes(hash_or_entity, raw)
  unless raw.empty?
    update_persistence_state(hash_or_entity, raw)
    key_condition = prepare_key_condition_from_entity(hash_or_entity)
    where(key_condition).__original_update(raw)
  end
  hash_or_entity
end

def update_persistence_state(entity, raw)
  persistence_state = entity.instance_variable_get(PERSISTENCE_STATE_VARIABLE)
  if persistence_state
    persistence_state.merge!(raw)
  end
end

def with_lock
  for_update
end

def with_mapper(mapper)
  clone(
    schema_mapper: mapper || opts[:schema_mapper],
  ).with_row_proc(
    Proc.new { |hash|
      m = mapper || opts[:schema_mapper]
      if m
        entity = m.db_load(hash)
        self.set_persistence_state(entity, hash)
        entity
      else
        hash
      end
    }
  )
end