class StatelyDB::Transaction::Transaction

def abort
  req = Stately::Db::TransactionRequest.new(
    abort: Google::Protobuf::Empty.new
  )
  resp = request_only(req)
  @is_transaction_open = false
  resp
end

def begin
  @is_transaction_open = true
  req = Stately::Db::TransactionRequest.new(
    begin: Stately::Db::TransactionBegin.new(store_id: @store_id.to_i,
                                             schema_id: @schema::SCHEMA_ID,
                                             schema_version_id: @schema::SCHEMA_VERSION_ID)
  )
  request_only(req)
  @incoming_responses = @stub.transaction(@outgoing_requests)
  nil
end

def begin_list(prefix,
               limit: 100,
               sort_property: nil,
               sort_direction: :ascending)
  sort_direction = case sort_direction
                   when :ascending
                     0
                   else
                     1
                   end
  req = Stately::Db::TransactionRequest.new(
    begin_list: Stately::Db::TransactionBeginList.new(
      key_path_prefix: String(prefix),
      limit:,
      sort_property:,
      sort_direction:
    )
  )
  do_list_request_response(req)
end

def commit
  req = Stately::Db::TransactionRequest.new(
    commit: Google::Protobuf::Empty.new
  )
  resp = request_response(req).finished
  @is_transaction_open = false
  Result.new(
    puts: resp.put_results.map do |result|
      @schema.unmarshal_item(stately_item: result)
    end,
    deletes: resp.delete_results.map(&:key_path)
  )
end

def continue_list(token, continue_direction: :forward)
  continue_direction = continue_direction == :forward ? 0 : 1
  req = Stately::Db::TransactionRequest.new(
    continue_list: Stately::Db::TransactionContinueList.new(
      token_data: token.token_data,
      direction: continue_direction
    )
  )
  do_list_request_response(req)
end

def delete(*key_paths)
  key_paths = Array(key_paths).flatten
  req = Stately::Db::TransactionRequest.new(
    delete_items: Stately::Db::TransactionDelete.new(
      deletes: key_paths.map { |key_path| Stately::Db::DeleteItem.new(key_path: String(key_path)) }
    )
  )
  request_only(req)
  nil
end

def do_list_request_response(req)
  items = []
  token = request_list_responses(req) do |resp|
    resp.result.items.each do |list_items_result|
      items << @schema.unmarshal_item(stately_item: list_items_result)
    end
  end
  [items, token]
end

def get(key_path)
  resp = get_batch(key_path)
  # Always return a single Item.
  resp.first
end

def get_batch(*key_paths)
  key_paths = Array(key_paths).flatten
  req = Stately::Db::TransactionRequest.new(
    get_items: Stately::Db::TransactionGet.new(
      gets: key_paths.map { |key_path| Stately::Db::GetItem.new(key_path: String(key_path)) }
    )
  )
  resp = request_response(req).get_results
  resp.items.map do |result|
    @schema.unmarshal_item(stately_item: result)
  end
end

def infer_response_type_from_request(req)
  if req.respond_to?(:get_items)
    Stately::Db::TransactionGetResponse
  elsif req.respond_to?(:list_items)
    Stately::Db::TransactionListResponse
  else
    raise "Unknown request type or request type does not have a corresponding response type"
  end
end

def infer_response_type_from_response(resp)
  if resp.respond_to?(:get_results)
    Stately::Db::TransactionGetResponse
  elsif resp.respond_to?(:list_results)
    Stately::Db::TransactionListResponse
  else
    raise "Unknown response type"
  end
end

def initialize(stub:, store_id:, schema:)
  @stub = stub
  @store_id = store_id
  @schema = schema
  @is_transaction_open = false
  # A queue of outbound requests
  @outgoing_requests = StatelyDB::Transaction::Queue.new
end

def open?
  @is_transaction_open
end

def put(item,
        must_not_exist: false,
        overwrite_metadata_timestamps: false)
  resp = put_batch({ item:, must_not_exist:, overwrite_metadata_timestamps: })
  resp.first
end

def put_batch(*items)
  puts = Array(items).flatten.map do |input|
    if input.is_a?(Hash)
      item = input[:item]
      Stately::Db::PutItem.new(
        item: item.send("marshal_stately"),
        overwrite_metadata_timestamps: input[:overwrite_metadata_timestamps],
        must_not_exist: input[:must_not_exist]
      )
    else
      Stately::Db::PutItem.new(
        item: input.send("marshal_stately")
      )
    end
  end
  req = Stately::Db::TransactionRequest.new(
    put_items: Stately::Db::TransactionPut.new(
      puts:
    )
  )
  resp = request_response(req).put_ack
  resp.generated_ids.map do |generated_id|
    case generated_id.value
    when :bytes
      StatelyDB::UUID.valid_uuid?(generated_id.bytes) ? StatelyDB::UUID.parse(generated_id.bytes) : generated_id.bytes
    when :uint
      generated_id.uint
    else # rubocop:disable Style/EmptyElse
      # An empty identifier is sent in the transaction Put response if an initialValue is not set
      nil
    end
  end
end

def request_list_responses(req)
  request_only(req)
  token = nil
  loop do
    resp = @incoming_responses.next.list_results
    if resp.finished
      raw_token = resp.finished.token
      token = StatelyDB::Token.new(token_data: raw_token.token_data,
                                   can_continue: raw_token.can_continue,
                                   can_sync: raw_token.can_sync,
                                   schema_version_id: raw_token.schema_version_id)
      break
    end
    yield resp
  end
  token
end

def request_only(req)
  req.message_id = @outgoing_requests.next_message_id
  @outgoing_requests.push(req)
  nil
end

def request_response(req)
  request_only(req)
  begin
    resp = @incoming_responses.next
    if req.message_id != resp.message_id
      raise "Message ID mismatch: request #{req.message_id} != response #{resp.message_id}"
    end
    raise "Response type mismatch" if infer_response_type_from_request(req) != infer_response_type_from_response(resp)
    resp
  rescue StopIteration
    nil
  end
end