// Copyright 2017 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package client

import (
	"sync"
	"time"

	"github.com/molecula/featurebase/v2/client/egpool"
	"github.com/molecula/featurebase/v2/logger"
	"github.com/molecula/featurebase/v2/roaring"
	"github.com/pkg/errors"
)

// Batch defaults.
const (
	DefaultKeyTranslateBatchSize = 100000
)

// TODO if using column translation, column ids might get way out of
// order. Could be worth sorting everything after translation (as an
// option?). Instead of sorting all simultaneously, it might be faster
// (more cache friendly) to sort ids and save the swap ops to apply to
// everything else that needs to be sorted.

// TODO support clearing values? nil values in records are ignored,
// but perhaps we could have a special type indicating that a bit or
// value should explicitly be cleared?

// RecordBatch is a Pilosa ingest interface designed to allow for
// maximum throughput on common workloads. Users should call Add()
// with a Row object until it returns ErrBatchNowFull, at which time
// they should call Import(), and then repeat.
//
// Add will not modify or otherwise retain the Row once it returns, so
// it is recommended that callers reuse the same Row with repeated
// calls to Add, just modifying its values appropriately in between
// calls. This avoids allocating a new slice of Values for each
// inserted Row.
//
// The supported types of the values in Row.Values are implementation
// defined. Similarly, the supported types for Row.ID are
// implementation defined.
type RecordBatch interface {
	Add(Row) error

	// Import does translation, creates the fragment files, and then,
	// if we're not using split batch mode, imports everything to
	// Pilosa. It then resets internal data structures for the next
	// batch. If we are using split batch mode, it saves the fragment
	// data to the batch, resets all other internal structures, and
	// continues.
	// Split batch mode DOES NOT CURRENTLY SUPPORT MUTEX OR INT FIELDS!
	Import() error

	// Len reports the number of records which have been added to the
	// batch since the last call to Import (or since it was created).
	Len() int

	// Flush is only applicable in split batch mode where it actually
	// imports the stored data to Pilosa. Otherwise it simply returns
	// nil.
	Flush() error
}

// agedTranslation combines a translation with a recording of when it was last used.
type agedTranslation struct {
	id       uint64
	lastUsed uint64
}

// Batch implements RecordBatch.
//
// It supports Values of type string, uint64, int64, or nil. The
// following table describes what Pilosa field each type of value must
// map to. Fields are set up when calling "NewBatch".
//
// | type   | pilosa field type | options   |
// |--------+-------------------+-----------|
// | string | set               | keys=true |
// | uint64 | set               | any       |
// | int64  | int               | any       |
// | float64| decimal           | scale     |
// | nil    | any               |           |
//
// nil values are ignored.
type Batch struct {
	client    *Client
	index     *Index
	header    []*Field
	headerMap map[string]*Field

	// prevDuration records the time that each doImport() takes. This
	// is used to set the timeout for transactions to a reasonable
	// value based on the last import. It starts with a conservative
	// default set in NewBatch.
	prevDuration time.Duration

	// ids is a slice of length batchSize of record IDs
	ids []uint64

	// rowIDs is a map of field index (in the header) to slices of
	// length batchSize which contain row IDs.
	rowIDs map[int][]uint64
	// clearRowIDs is a map[fieldIndex][idsIndex]rowID we don't expect
	// clears to happen very often, so we store the idIndex/value
	// mapping in a map rather than a slice as we do for rowIDs. This
	// is a potentially temporary workaround to allow packed boolean
	// fields to clear "false" values. Packed fields may be more
	// completely supported by Pilosa in future.
	clearRowIDs map[int]map[int]uint64

	// rowIDSets is a map from field name to a batchSize slice of
	// slices of row IDs. When a given record can have more than one
	// value for a field, rowIDSets stores that information.
	rowIDSets map[string][][]uint64

	// values holds the values for each record of an int field
	values map[string][]int64

	// times holds a time for each record. (if any of the fields are time fields)
	times []QuantizedTime

	// nullIndices holds a slice of indices into b.ids for each
	// integer field which has nil values.
	nullIndices map[string][]uint64

	// TODO support bool fields.

	// for each field, keep a map of key to which record indexes that key mapped to
	toTranslate      map[int]map[string][]int
	toTranslateClear map[int]map[string][]int

	// toTranslateSets is a map from field name to a map of string
	// keys that need to be translated to sets of record indexes which
	// those keys map to.
	toTranslateSets map[string]map[string][]int

	// toTranslateID maps each string key to a record index - this
	// will get translated into Batch.rowIDs
	toTranslateID []string

	colTranslations map[string]agedTranslation
	rowTranslations map[string]map[string]agedTranslation
	cycle           uint64
	maxAge          uint64

	// staleTime tracks the time the first record of the batch was inserted
	// plus the maxStaleness, in order to raise ErrBatchNowStale if the
	// maxStaleness has elapsed
	staleTime    time.Time
	maxStaleness time.Duration

	// Maximum number of keys to translate at one time.
	keyTranslateBatchSize int

	log logger.Logger

	// experimental — only used by FlushToFragments which is an
	// alternative to Import which just builds the bitmap data for a
	// batch without actually importing it.
	splitBatchMode bool
	frags          fragments
	clearFrags     fragments
}

func (b *Batch) Len() int { return len(b.ids) }

// BatchOption is a functional option for Batch objects.
type BatchOption func(b *Batch) error

func OptLogger(l logger.Logger) BatchOption {
	return func(b *Batch) error {
		b.log = l
		return nil
	}
}

func OptSplitBatchMode(on bool) BatchOption {
	return func(b *Batch) error {
		b.splitBatchMode = on
		return nil
	}
}

func OptCacheMaxAge(age uint64) BatchOption {
	return func(b *Batch) error {
		b.maxAge = age
		return nil
	}
}

func OptMaxStaleness(t time.Duration) BatchOption {
	return func(b *Batch) error {
		b.maxStaleness = t
		return nil
	}
}

func OptKeyTranslateBatchSize(v int) BatchOption {
	return func(b *Batch) error {
		b.keyTranslateBatchSize = v
		return nil
	}
}

// NewBatch initializes a new Batch object which will use the given
// Pilosa client, index, set of fields, and will take "size" records
// before returning ErrBatchNowFull. The positions of the Fields in
// 'fields' correspond to the positions of values in the Row's Values
// passed to Batch.Add().
func NewBatch(client *Client, size int, index *Index, fields []*Field, opts ...BatchOption) (*Batch, error) {
	if len(fields) == 0 || size == 0 {
		return nil, errors.New("can't batch with no fields or batch size")
	}
	headerMap := make(map[string]*Field, len(fields))
	rowIDs := make(map[int][]uint64, len(fields))
	values := make(map[string][]int64)
	tt := make(map[int]map[string][]int, len(fields))
	ttSets := make(map[string]map[string][]int)
	hasTime := false
	for i, field := range fields {
		headerMap[field.Name()] = field
		opts := field.Opts()
		switch typ := opts.Type(); typ {
		case FieldTypeDefault, FieldTypeSet, FieldTypeTime:
			if opts.Keys() {
				tt[i] = make(map[string][]int)
				ttSets[field.Name()] = make(map[string][]int)
			}
			hasTime = typ == FieldTypeTime || hasTime
		case FieldTypeInt, FieldTypeDecimal, FieldTypeTimestamp:
			// tt line only needed if int field is string foreign key
			tt[i] = make(map[string][]int)
			values[field.Name()] = make([]int64, 0, size)
		case FieldTypeMutex:
			// similar to set/time fields, but no need to support sets
			// of values (hence no ttSets)
			if opts.Keys() {
				tt[i] = make(map[string][]int)
			}
			rowIDs[i] = make([]uint64, 0, size)
		default:
			return nil, errors.Errorf("field type '%s' is not currently supported through Batch", typ)
		}
	}

	b := &Batch{
		client:                client,
		header:                fields,
		headerMap:             headerMap,
		prevDuration:          time.Minute * 11,
		index:                 index,
		ids:                   make([]uint64, 0, size),
		rowIDs:                rowIDs,
		clearRowIDs:           make(map[int]map[int]uint64),
		rowIDSets:             make(map[string][][]uint64),
		values:                values,
		nullIndices:           make(map[string][]uint64),
		toTranslate:           tt,
		toTranslateClear:      make(map[int]map[string][]int),
		toTranslateSets:       ttSets,
		colTranslations:       make(map[string]agedTranslation),
		rowTranslations:       make(map[string]map[string]agedTranslation),
		maxAge:                64,
		maxStaleness:          time.Duration(0),
		keyTranslateBatchSize: DefaultKeyTranslateBatchSize,

		log: logger.NopLogger,

		frags:      make(fragments),
		clearFrags: make(fragments),
	}
	if hasTime {
		b.times = make([]QuantizedTime, 0, size)
	}
	for _, opt := range opts {
		err := opt(b)
		if err != nil {
			return nil, errors.Wrap(err, "applying options")
		}
	}
	return b, nil
}

// Row represents a single record which can be added to a Batch.
type Row struct {
	ID interface{}
	// Values map to the slice of fields in Batch.header
	Values []interface{}
	// Clears' int key is an index into Batch.header
	Clears map[int]interface{}
	// Time applies to all time fields
	Time QuantizedTime
}

// QuantizedTime represents a moment in time down to some granularity
// (year, month, day, or hour).
type QuantizedTime struct {
	ymdh [10]byte
}

// Set sets the Quantized time to the given timestamp (down to hour
// granularity).
func (qt *QuantizedTime) Set(t time.Time) {
	copy(qt.ymdh[:], t.Format("2006010215"))
}

// SetYear sets the quantized time's year, but leaves month, day, and
// hour untouched.
func (qt *QuantizedTime) SetYear(year string) {
	copy(qt.ymdh[:4], year)
}

// SetMonth sets the QuantizedTime's month, but leaves year, day, and
// hour untouched.
func (qt *QuantizedTime) SetMonth(month string) {
	copy(qt.ymdh[4:6], month)
}

// SetDay sets the QuantizedTime's day, but leaves year, month, and
// hour untouched.
func (qt *QuantizedTime) SetDay(day string) {
	copy(qt.ymdh[6:8], day)
}

// SetHour sets the QuantizedTime's hour, but leaves year, month, and
// day untouched.
func (qt *QuantizedTime) SetHour(hour string) {
	copy(qt.ymdh[8:10], hour)
}

// Reset sets the time to the zero value which generates no time views.
func (qt *QuantizedTime) Reset() {
	for i := range qt.ymdh {
		qt.ymdh[i] = 0
	}
}

// views builds the list of Pilosa views for this particular time,
// given a quantum.
func (qt *QuantizedTime) views(q TimeQuantum) ([]string, error) {
	zero := QuantizedTime{}
	if *qt == zero {
		return nil, nil
	}
	views := make([]string, 0, len(q))
	for _, unit := range q {
		switch unit {
		case 'Y':
			if qt.ymdh[0] == 0 {
				return nil, errors.New("no data set for year")
			}
			views = append(views, string(qt.ymdh[:4]))
		case 'M':
			if qt.ymdh[4] == 0 {
				return nil, errors.New("no data set for month")
			}
			views = append(views, string(qt.ymdh[:6]))
		case 'D':
			if qt.ymdh[6] == 0 {
				return nil, errors.New("no data set for day")
			}
			views = append(views, string(qt.ymdh[:8]))
		case 'H':
			if qt.ymdh[8] == 0 {
				return nil, errors.New("no data set for hour")
			}
			views = append(views, string(qt.ymdh[:10]))
		}
	}
	return views, nil
}

func (b *Batch) getColTranslation(key string) (uint64, bool) {
	trans, ok := b.colTranslations[key]
	if ok {
		trans.lastUsed = b.cycle
		b.colTranslations[key] = trans
	}
	return trans.id, ok
}

func (b *Batch) getRowTranslation(field, key string) (uint64, bool) {
	trans, ok := b.rowTranslations[field][key]
	if ok {
		trans.lastUsed = b.cycle
		b.rowTranslations[field][key] = trans
	}
	return trans.id, ok
}

// Add adds a record to the batch. Performance will be best if record
// IDs are shard-sorted. That is, all records which belong to the same
// Pilosa shard are added adjacent to each other. If the records are
// also in-order within a shard this will likely help as well. Add
// clears rec.Clears when it returns normally (either a nil error or
// BatchNowFull).
func (b *Batch) Add(rec Row) error {
	// Clear recValues and rec.Clears upon return.
	defer func() {
		for i := range rec.Values {
			rec.Values[i] = nil
		}
		for k := range rec.Clears {
			delete(rec.Clears, k)
		}
	}()

	if len(b.ids) == cap(b.ids) {
		return ErrBatchAlreadyFull
	}
	if len(rec.Values) != len(b.header) {
		return errors.Errorf("record needs to match up with batch fields, got %d fields and %d record", len(b.header), len(rec.Values))
	}

	handleStringID := func(rid string) error {
		if rid == "" {
			return errors.Errorf("record identifier cannot be an empty string")
		}
		if colID, ok := b.getColTranslation(rid); ok {
			b.ids = append(b.ids, colID)
		} else {
			if b.toTranslateID == nil {
				b.toTranslateID = make([]string, cap(b.ids))
			}
			b.toTranslateID[len(b.ids)] = rid
			b.ids = append(b.ids, 0)
		}
		return nil
	}
	var err error

	switch rid := rec.ID.(type) {
	case uint64:
		b.ids = append(b.ids, rid)
	case string:
		err := handleStringID(rid)
		if err != nil {
			return err
		}
	case []byte:
		err = handleStringID(string(rid))
		if err != nil {
			return err
		}
	default: // TODO support nil ID as being auto-allocated.
		return errors.Errorf("unsupported id type %T value %v", rid, rid)
	}

	// curPos is the current position in b.ids, rowIDs[*], etc.
	curPos := len(b.ids) - 1

	if b.times != nil {
		b.times = append(b.times, rec.Time)
	}

	for i := 0; i < len(rec.Values); i++ {
		field := b.header[i]
		switch val := rec.Values[i].(type) {
		case string:
			if field.Opts().Type() != FieldTypeInt {
				// nil-extend
				for len(b.rowIDs[i]) < curPos {
					b.rowIDs[i] = append(b.rowIDs[i], nilSentinel)
				}
				rowIDs := b.rowIDs[i]
				// empty string is not a valid value at this point (Pilosa refuses to translate it)
				if val == "" { //
					b.rowIDs[i] = append(rowIDs, nilSentinel)

				} else if rowID, ok := b.getRowTranslation(field.Name(), val); ok {
					b.rowIDs[i] = append(rowIDs, rowID)
				} else {
					ints, ok := b.toTranslate[i][val]
					if !ok {
						ints = make([]int, 0)
					}
					ints = append(ints, curPos)
					b.toTranslate[i][val] = ints
					b.rowIDs[i] = append(rowIDs, 0)
				}
			} else if field.Opts().Type() == FieldTypeInt {
				if val == "" {
					// copied from the `case nil:` section for ints and decimals
					b.values[field.Name()] = append(b.values[field.Name()], 0)
					nullIndices, ok := b.nullIndices[field.Name()]
					if !ok {
						nullIndices = make([]uint64, 0)
					}
					nullIndices = append(nullIndices, uint64(curPos))
					b.nullIndices[field.Name()] = nullIndices
				} else if intVal, ok := b.getRowTranslation(field.Name(), val); ok {
					b.values[field.Name()] = append(b.values[field.Name()], int64(intVal))
				} else {
					ints, ok := b.toTranslate[i][val]
					if !ok {
						ints = make([]int, 0)
					}
					ints = append(ints, curPos)
					b.toTranslate[i][val] = ints
					b.values[field.Name()] = append(b.values[field.Name()], 0)
				}
			}
		case uint64:
			// nil-extend
			for len(b.rowIDs[i]) < curPos {
				b.rowIDs[i] = append(b.rowIDs[i], nilSentinel)
			}
			b.rowIDs[i] = append(b.rowIDs[i], val)
		case int64:
			b.values[field.Name()] = append(b.values[field.Name()], val)
		case []string:
			if len(val) == 0 {
				continue
			}
			rowIDSets, ok := b.rowIDSets[field.Name()]
			if !ok {
				rowIDSets = make([][]uint64, len(b.ids)-1, cap(b.ids))
				b.rowIDSets[field.Name()] = rowIDSets
			}
			for len(rowIDSets) < len(b.ids)-1 {
				rowIDSets = append(rowIDSets, nil) // nil extend
			}

			rowIDs := make([]uint64, 0, len(val))
			for _, k := range val {
				if k == "" {
					continue
				}
				if rowID, ok := b.getRowTranslation(field.Name(), k); ok {
					rowIDs = append(rowIDs, rowID)
				} else {
					ttsets, ok := b.toTranslateSets[field.Name()]
					if !ok {
						ttsets = make(map[string][]int)
						b.toTranslateSets[field.Name()] = make(map[string][]int)
					}
					ints, ok := ttsets[k]
					if !ok {
						ints = make([]int, 0, 1)
					}
					ints = append(ints, curPos)
					b.toTranslateSets[field.Name()][k] = ints
				}
			}
			b.rowIDSets[field.Name()] = append(rowIDSets, rowIDs)
		case []uint64:
			if len(val) == 0 {
				continue
			}
			rowIDSets, ok := b.rowIDSets[field.Name()]
			if !ok {
				rowIDSets = make([][]uint64, len(b.ids)-1, cap(b.ids))
			}
			for len(rowIDSets) < len(b.ids)-1 {
				rowIDSets = append(rowIDSets, nil) // nil extend
			}
			b.rowIDSets[field.Name()] = append(rowIDSets, val)
		case nil:
			if field.Opts().Type() == FieldTypeInt || field.Opts().Type() == FieldTypeDecimal {
				b.values[field.Name()] = append(b.values[field.Name()], 0)
				nullIndices, ok := b.nullIndices[field.Name()]
				if !ok {
					nullIndices = make([]uint64, 0)
				}
				nullIndices = append(nullIndices, uint64(curPos))
				b.nullIndices[field.Name()] = nullIndices

			} else {
				// only append nil to rowIDs if this field already has
				// rowIDs. Otherwise, this could be a []string or
				// []uint64 field where we've only seen nil values so
				// far. when we see a uint64 or string value, we'll
				// "nil-extend" rowIDs to make sure it's the right
				// length.
				if rowIDs, ok := b.rowIDs[i]; ok {
					b.rowIDs[i] = append(rowIDs, nilSentinel)
				}
			}
		default:
			return errors.Errorf("Val %v Type %[1]T is not currently supported. Use string, uint64 (row id), or int64 (integer value)", val)
		}
	}

	for i, uval := range rec.Clears {
		field := b.header[i]
		if _, ok := b.clearRowIDs[i]; !ok {
			b.clearRowIDs[i] = make(map[int]uint64)
		}
		switch val := uval.(type) {
		case string:
			clearRows := b.clearRowIDs[i]
			// translate val and add to clearRows
			if rowID, ok := b.getRowTranslation(field.Name(), val); ok {
				clearRows[curPos] = rowID
			} else {
				_, ok := b.toTranslateClear[i]
				if !ok {
					b.toTranslateClear[i] = make(map[string][]int)
				}
				ints, ok := b.toTranslateClear[i][val]
				if !ok {
					ints = make([]int, 0)
				}
				ints = append(ints, curPos)
				b.toTranslateClear[i][val] = ints
			}
		case uint64:
			b.clearRowIDs[i][curPos] = val
		default:
			return errors.Errorf("Clearing a value '%v' Type %[1]T is not currently supported (field '%s')", val, field.Name())
		}
		// nil extend b.rowIDs so we don't run into a horrible bug
		// where we skip doing clears because b.rowIDs doesn't have a
		// value for this field
		for len(b.rowIDs[i]) <= curPos {
			b.rowIDs[i] = append(b.rowIDs[i], nilSentinel)
		}

	}

	if len(b.ids) == cap(b.ids) {
		return ErrBatchNowFull
	}
	if b.maxStaleness != time.Duration(0) { // set maxStaleness to 0 to disable staleness checking
		if len(b.ids) == 1 {
			b.staleTime = time.Now().Add(b.maxStaleness)
		} else if time.Now().After(b.staleTime) {
			return ErrBatchNowStale
		}
	}
	return nil
}

// ErrBatchNowFull, similar to io.EOF, is a marker error to notify the
// user of a batch that it is time to call Import.
var ErrBatchNowFull = errors.New("batch is now full - you cannot add any more records (though the one you just added was accepted)")

// ErrBatchAlreadyFull is a real error saying that Batch.Add did not
// complete because the batch was full.
var ErrBatchAlreadyFull = errors.New("batch was already full, record was rejected")

// ErrBatchNowStale indicates that the oldest record in the batch is older than
// the maxStaleness value of the batch. Like ErrBatchNowFull, the error does
// not mean the record was rejected.
var ErrBatchNowStale = errors.New("batch is stale and needs to be imported (however, record was accepted)")

// Import does translation, creates the fragment files, and then,
// if we're not using split batch mode, imports everything to
// Pilosa. It then resets internal data structures for the next
// batch. If we are using split batch mode, it saves the fragment
// data to the batch, resets all other internal structures, and
// continues.  split batch mode DOES NOT CURRENTLY SUPPORT MUTEX
// OR INT FIELDS!
func (b *Batch) Import() error {
	start := time.Now()
	trns, err := b.client.StartTransaction("", b.prevDuration*10, false, time.Hour)
	if err != nil {
		return errors.Wrap(err, "starting transaction")
	}
	defer func() {
		trns, err := b.client.FinishTransaction(trns.ID)
		if err != nil {
			b.log.Errorf("error finishing transaction: %v. trns: %+v", err, trns)
		}
		b.client.Stats.Timing(MetricBatchImportDurationSeconds, time.Since(start), 1.0)
	}()

	size := len(b.ids)
	transStart := time.Now()
	// first we need to translate the toTranslate, then fill out the missing row IDs
	err = b.doTranslation()
	if err != nil {
		return errors.Wrap(err, "doing Translation")
	}
	transTime := time.Now()
	b.log.Printf("translating batch of %d took: %v", size, transTime.Sub(transStart))

	frags, clearFrags, err := b.makeFragments(b.frags, b.clearFrags)
	if err != nil {
		return errors.Wrap(err, "making fragments (flush)")
	}
	makeTime := time.Now()
	b.log.Printf("making fragments for batch of %d took %v", size, makeTime.Sub(transTime))

	if b.splitBatchMode {
		b.frags = frags
		b.clearFrags = clearFrags
	} else {
		b.frags = make(fragments)
		b.clearFrags = make(fragments)
		// create bitmaps out of each field in b.rowIDs and import. Also
		// import int data.
		err = b.doImport(frags, clearFrags)
		if err != nil {
			return errors.Wrap(err, "doing import")
		}
		b.log.Printf("importing fragments took %v", time.Since(makeTime))
	}

	b.reset()
	return nil
}

// Flush is only applicable in split batch mode where it actually
// imports the stored data to Pilosa. Otherwise it simply returns
// nil.
func (b *Batch) Flush() error {
	if !b.splitBatchMode {
		return nil
	}
	start := time.Now()

	trns, err := b.client.StartTransaction("", b.prevDuration*10, false, time.Hour)
	if err != nil {
		return errors.Wrap(err, "starting transaction")
	}
	defer func() {
		trns, err := b.client.FinishTransaction(trns.ID)
		if err != nil {
			b.log.Errorf("error finishing transaction: %v. trns: %+v", err, trns)
		}
		b.client.Stats.Timing(MetricBatchFlushDurationSeconds, time.Since(start), 1.0)
	}()

	importStart := time.Now()
	err = b.doImport(b.frags, b.clearFrags)
	if err != nil {
		return errors.Wrap(err, "doing import (ImportFragments)")
	}

	b.log.Debugf("superbatch import took %v", time.Since(importStart))

	b.reset()
	b.frags = make(fragments)
	b.clearFrags = make(fragments)
	return nil
}

func (b *Batch) doTranslation() error {
	eg := egpool.Group{PoolSize: 20}

	// Translate the column keys.
	eg.Go(func() error {
		// Dedupliucate keys to translate.
		dedup := make(map[string]struct{})
		var keys []string
		for _, key := range b.toTranslateID {
			if key == "" {
				continue
			}

			if _, ok := dedup[key]; ok {
				continue
			}
			dedup[key] = struct{}{}

			keys = append(keys, key)
		}
		if len(keys) == 0 {
			// There are no column keys to translate.
			return nil
		}

		// Create the keys.
		start := time.Now()
		trans, err := b.createIndexKeys(b.index, keys...)
		if err != nil {
			return errors.Wrap(err, "translating col keys")
		}
		if len(trans) != len(keys) {
			return errors.Errorf("requested IDs for %d column keys but got %d back", len(keys), len(trans))
		}
		b.log.Debugf("translating %d column keys took %v", len(keys), time.Since(start))

		// Apply keys to translation cache.
		for key, id := range trans {
			b.colTranslations[key] = agedTranslation{
				id:       id,
				lastUsed: b.cycle,
			}
		}

		// Translate remaining keys in batch.
		for index, ttkey := range b.toTranslateID {
			if ttkey == "" {
				continue
			}

			b.ids[index] = trans[ttkey]
		}

		return nil
	})

	// creating a lock up here for the rowCache(s) which we get
	// below. Usually this isn't needed, but sometimes I think the
	// same rowCache gets used repeatedly because the same field is in
	// there multiple times, and that can lead to race
	// conditions. Need to understand this better, but gonna see if
	// this avoids the races.
	rowCacheLock := &sync.Mutex{}

	// Translate the row keys.
	for i, tt := range b.toTranslate {
		// Skip this if there are no keys to translate.
		ttc := b.toTranslateClear[i]
		if len(tt) == 0 && len(ttc) == 0 {
			continue
		}

		// Look up the associated field.
		field := b.header[i]
		fieldName := field.Name()

		// Fetch the translation cache.
		rowCache := b.rowTranslations[fieldName]
		if rowCache == nil {
			rowCache = make(map[string]agedTranslation)
			b.rowTranslations[fieldName] = rowCache
		}

		i, tt := i, tt
		eg.Go(func() error {
			// Collect the keys to translate.
			keys := make([]string, 0, len(tt)+len(ttc))
			for k := range tt {
				keys = append(keys, k)
			}
			for k := range ttc {
				keys = append(keys, k)
			}

			// Create the keys.
			start := time.Now()
			trans, err := b.createFieldKeys(field, keys...)
			if err != nil {
				return errors.Wrap(err, "translating field keys")
			}
			b.log.Debugf("translating %d field keys for %s took %v", len(trans), fieldName, time.Since(start))

			// Apply keys to translation cache.
			rowCacheLock.Lock()
			for key, id := range trans {
				rowCache[key] = agedTranslation{
					id:       id,
					lastUsed: b.cycle,
				}
			}
			rowCacheLock.Unlock()

			switch ftype := field.Opts().Type(); ftype {
			case FieldTypeSet, FieldTypeMutex, FieldTypeTime:
				// Fill out missing IDs in local batch records with translated IDs.
				rows := b.rowIDs[i]
				for key, idxs := range tt {
					id, ok := trans[key]
					if !ok {
						return errors.Errorf("key translation missing: %q in field %q", key, fieldName)
					}

					for _, i := range idxs {
						rows[i] = id
					}
				}

				// Fill out missing IDs in clear lists.
				clearRows := b.clearRowIDs[i]
				for key, idxs := range ttc {
					id, ok := trans[key]
					if !ok {
						return errors.Errorf("key translation missing: %q in field %q", key, fieldName)
					}

					for _, i := range idxs {
						clearRows[i] = id
					}
				}

			case FieldTypeInt:
				// Handle foreign key int fields — fill out b.values instead of b.rows.
				vals := b.values[fieldName]
				for key, idxs := range tt {
					id, ok := trans[key]
					if !ok {
						return errors.Errorf("key translation missing: %q in field %q", key, fieldName)
					}

					for _, i := range idxs {
						vals[i] = int64(id)
					}
				}

			default:
				return errors.Errorf("unexpected field type for translation: %q", ftype)
			}

			return nil
		})
	}

	for fieldName, tt := range b.toTranslateSets {
		// Skip this if there are no keys to translate.
		if len(tt) == 0 {
			continue
		}

		// Look up the associated field.
		field := b.headerMap[fieldName]

		// Fetch the translation cache.
		rowCache := b.rowTranslations[fieldName]
		if rowCache == nil {
			rowCache = make(map[string]agedTranslation)
			b.rowTranslations[fieldName] = rowCache
		}

		fieldName, tt := fieldName, tt
		eg.Go(func() error {
			// Collect the keys to translate.
			keys := make([]string, 0, len(tt))
			for k := range tt {
				keys = append(keys, k)
			}

			// Create the keys.
			start := time.Now()
			trans, err := b.createFieldKeys(field, keys...)
			if err != nil {
				return errors.Wrap(err, "translating field keys")
			}
			b.log.Debugf("translating %d field keys for %s took %v", len(trans), fieldName, time.Since(start))

			// Apply keys to translation cache.
			rowCacheLock.Lock()
			for key, id := range trans {
				rowCache[key] = agedTranslation{
					id:       id,
					lastUsed: b.cycle,
				}
			}
			rowCacheLock.Unlock()

			// Fill out missing IDs in local batch records with translated IDs.
			rowIDSets := b.rowIDSets[fieldName]
			for key, idxs := range tt {
				id, ok := trans[key]
				if !ok {
					return errors.Errorf("key translation missing: %q in field %q", key, fieldName)
				}

				for _, i := range idxs {
					rowIDSets[i] = append(rowIDSets[i], id)
				}
			}

			return nil
		})
	}

	return eg.Wait()
}

func (b *Batch) createIndexKeys(index *Index, keys ...string) (map[string]uint64, error) {
	batchSize := b.keyTranslateBatchSize
	if batchSize <= 0 || len(keys) <= batchSize {
		return b.client.CreateIndexKeys(index, keys...)
	}

	results := make(map[string]uint64, len(keys))
	for len(keys) > 0 {
		keySlice := keys
		if len(keySlice) > batchSize {
			keySlice = keySlice[:batchSize]
		}

		trans, err := b.client.CreateIndexKeys(index, keySlice...)
		if err != nil {
			return nil, err
		} else if len(trans) != len(keySlice) {
			return nil, errors.Errorf("requested IDs for %d column keys but got %d back", len(keySlice), len(trans))
		}
		for key, id := range trans {
			results[key] = id
		}

		keys = keys[len(keySlice):]
	}

	return results, nil
}

func (b *Batch) createFieldKeys(field *Field, keys ...string) (map[string]uint64, error) {
	batchSize := b.keyTranslateBatchSize
	if batchSize <= 0 || len(keys) <= batchSize {
		return b.client.CreateFieldKeys(field, keys...)
	}

	results := make(map[string]uint64, len(keys))
	for len(keys) > 0 {
		keySlice := keys
		if len(keySlice) > batchSize {
			keySlice = keySlice[:batchSize]
		}

		trans, err := b.client.CreateFieldKeys(field, keySlice...)
		if err != nil {
			return nil, err
		} else if len(trans) != len(keySlice) {
			return nil, errors.Errorf("requested IDs for %d row keys but got %d back", len(keySlice), len(trans))
		}
		for key, id := range trans {
			results[key] = id
		}

		keys = keys[len(keySlice):]
	}

	return results, nil
}

func (b *Batch) doImport(frags, clearFrags fragments) error {

	start := time.Now()
	eg := egpool.Group{PoolSize: 20}
	// TODO, currently this relies on upstream behavior of
	// makeFragments to guarantee that any shard/field combination in
	// clearFrags also has a shard/field in frags. We're only
	// iterating over frags and then checking to see if clearFrags has
	// the same keys. If we optimized makeFragments to skip adding
	// things to frags which had no set bits (e.g. if we were only
	// clearing things), then this code would need to be updated to
	// ensure that it looked at the things in clearFrags which were
	// *not* in frags.
	for fragmentKey, viewMap := range frags {
		field := fragmentKey.field
		shard := fragmentKey.shard
		viewMap := viewMap

		eg.Go(func() error {
			clearViewMap := clearFrags.GetViewMap(shard, field)
			if len(clearViewMap) > 0 {
				start := time.Now()
				err := b.client.ImportRoaringBitmap(b.index.Field(field), shard, clearViewMap, true)
				if err != nil {
					return errors.Wrapf(err, "import clearing clearing data for %s", field)
				}
				b.log.Debugf("imp-roar-clr %s,shard:%d,views:%d %v", field, shard, len(clearViewMap), time.Since(start))
			}

			start := time.Now()
			err := b.client.ImportRoaringBitmap(b.index.Field(field), shard, viewMap, false)
			b.log.Debugf("imp-roar     %s,shard:%d,views:%d %v", field, shard, len(clearViewMap), time.Since(start))
			return errors.Wrapf(err, "importing data for %s", field)
		})
	}
	eg.Go(func() error { return b.importValueData() })
	eg.Go(func() error { return b.importMutexData() })

	err := eg.Wait()
	if err != nil {
		if pferr := anyCause(ErrPreconditionFailed, eg.Errors()...); pferr != nil {
			return pferr
		}
		return err
	}
	b.prevDuration = time.Since(start)
	return nil
}

func anyCause(cause error, errs ...error) error {
	if cause == nil {
		return nil
	}

	for _, err := range errs {
		if errors.Cause(err) == cause {
			return err
		}
	}
	return nil
}

// this is kind of bad as it means we can never import column id
// ^uint64(0) which is a valid column ID. I think it's unlikely to
// matter much in practice (we could maybe special case it somewhere
// if needed though).
var nilSentinel = ^uint64(0)

func (b *Batch) makeFragments(frags, clearFrags fragments) (fragments, fragments, error) {
	shardWidth := b.index.ShardWidth()
	if shardWidth == 0 {
		shardWidth = DefaultShardWidth
	}
	emptyClearRows := make(map[int]uint64)

	// create _exists fragments if needed
	if b.index.Opts().TrackExistence() {
		var curBM *roaring.Bitmap
		curShard := ^uint64(0) // impossible sentinel value for shard.
		for _, col := range b.ids {
			if col/shardWidth != curShard {
				curShard = col / shardWidth
				curBM = frags.GetOrCreate(curShard, "_exists", "")
			}
			curBM.DirectAdd(col % shardWidth)
		}
	}

	for i, rowIDs := range b.rowIDs {
		if len(rowIDs) == 0 {
			continue // this can happen when the values that came in for this field were string slices
		}
		clearRows := b.clearRowIDs[i]
		if clearRows == nil {
			clearRows = emptyClearRows
		}
		field := b.header[i]
		opts := field.Opts()
		if opts.Type() == FieldTypeMutex {
			continue // we handle mutex fields separately — they can't use importRoaring
		}
		curShard := ^uint64(0) // impossible sentinel value for shard.
		var curBM *roaring.Bitmap
		var clearBM *roaring.Bitmap
		for j := range b.ids {
			col, row := b.ids[j], rowIDs[j]
			if col/shardWidth != curShard {
				curShard = col / shardWidth
				curBM = frags.GetOrCreate(curShard, field.Name(), "")
				clearBM = clearFrags.GetOrCreate(curShard, field.Name(), "")
			}
			if row != nilSentinel {
				// TODO this is super ugly, but we want to avoid setting
				// bits on the standard view in the specific case when
				// there isn't one. Should probably refactor this whole
				// loop to be more general w.r.t. views. Also... tests for
				// the NoStandardView case would be great.
				if !(opts.Type() == FieldTypeTime && opts.NoStandardView()) {
					curBM.DirectAdd(row*shardWidth + (col % shardWidth))
				}
				if opts.Type() == FieldTypeTime {
					views, err := b.times[j].views(opts.TimeQuantum())
					if err != nil {
						return nil, nil, errors.Wrap(err, "calculating views")
					}
					for _, view := range views {
						tbm := frags.GetOrCreate(curShard, field.Name(), view)
						tbm.DirectAdd(row*shardWidth + (col % shardWidth))
					}
				}
			}

			clearRow, ok := clearRows[j]
			if ok {
				clearBM.DirectAddN(clearRow*shardWidth + (col % shardWidth))
				// we're going to execute the clear before the set, so
				// we want to make sure that at this point, the "set"
				// fragments don't contain the bit that we're clearing
				curBM.DirectRemoveN(clearRow*shardWidth + (col % shardWidth))
			}
		}
	}

	for fname, rowIDSets := range b.rowIDSets {
		if len(rowIDSets) == 0 {
			continue
		} else if len(rowIDSets) < len(b.ids) {
			// rowIDSets is guaranteed to have capacity == to b.ids,
			// but if the last record had a nil for this field, it
			// might not have the same length, so we re-slice it to
			// ensure the lengths are the same.
			rowIDSets = rowIDSets[:len(b.ids)]
		}
		field := b.headerMap[fname]
		opts := field.Opts()
		curShard := ^uint64(0) // impossible sentinel value for shard.
		var curBM *roaring.Bitmap
		for j := range b.ids {
			col, rowIDs := b.ids[j], rowIDSets[j]
			if len(rowIDs) == 0 {
				continue
			}
			if col/shardWidth != curShard {
				curShard = col / shardWidth
				curBM = frags.GetOrCreate(curShard, fname, "")
			}
			// TODO this is super ugly, but we want to avoid setting
			// bits on the standard view in the specific case when
			// there isn't one. Should probably refactor this whole
			// loop to be more general w.r.t. views. Also... tests for
			// the NoStandardView case would be great.
			if !(opts.Type() == FieldTypeTime && opts.NoStandardView()) {
				for _, row := range rowIDs {
					curBM.DirectAdd(row*shardWidth + (col % shardWidth))
				}
			}
			if opts.Type() == FieldTypeTime {
				views, err := b.times[j].views(opts.TimeQuantum())
				if err != nil {
					return nil, nil, errors.Wrap(err, "calculating views")
				}
				for _, view := range views {
					tbm := frags.GetOrCreate(curShard, fname, view)
					for _, row := range rowIDs {
						tbm.DirectAdd(row*shardWidth + (col % shardWidth))
					}
				}
			}
		}
	}
	return frags, clearFrags, nil
}

// importValueData imports data for int fields.
func (b *Batch) importValueData() error {
	shardWidth := b.index.ShardWidth()
	if shardWidth == 0 {
		shardWidth = DefaultShardWidth
	}
	eg := egpool.Group{PoolSize: 20}

	ids := make([]uint64, len(b.ids))
	for fieldName, bvalues := range b.values {
		ids = ids[:len(b.ids)]

		// trim out null values from ids and values.
		nullIndices := b.nullIndices[fieldName]

		i, n := uint64(0), 0
		for _, nullIndex := range nullIndices {
			copy(ids[n:], b.ids[i:nullIndex])
			n += copy(bvalues[n:], bvalues[i:nullIndex])
			i = nullIndex + 1
		}

		copy(ids[n:], b.ids[i:])
		n += copy(bvalues[n:], bvalues[i:])
		ids, bvalues = ids[:n], bvalues[:n]

		// now do imports by shard
		if len(ids) == 0 {
			continue // TODO test this "all nil" case
		}
		curShard := ids[0] / shardWidth
		startIdx := 0
		for i := 1; i <= len(ids); i++ {
			var recordID uint64
			if i < len(ids) {
				recordID = ids[i]
			} else {
				recordID = (curShard + 2) * shardWidth
			}

			if recordID/shardWidth != curShard {
				endIdx := i
				shard := curShard
				field := b.headerMap[fieldName]
				path, data, err := b.client.EncodeImportValues(field, shard, bvalues[startIdx:endIdx], ids[startIdx:endIdx], false)
				if err != nil {
					return errors.Wrap(err, "encoding import values")
				}
				eg.Go(func() error {
					start := time.Now()
					err := b.client.DoImportValues(b.index.Name(), shard, path, data)
					b.log.Debugf("imp-vals %s,shard:%d,data:%d %v", field, shard, len(data), time.Since(start))
					return errors.Wrapf(err, "importing values for %s", field)
				})
				startIdx = i
				curShard = recordID / shardWidth
			}
		}
	}
	err := eg.Wait()
	if err != nil {
		if pferr := anyCause(ErrPreconditionFailed, eg.Errors()...); pferr != nil {
			return pferr
		}
		return err
	}
	return errors.Wrap(err, "importing value data")
}

// TODO this should work for bools as well - just need to support them
// at batch creation time and when calling Add, I think.
func (b *Batch) importMutexData() error {
	shardWidth := b.index.ShardWidth()
	if shardWidth == 0 {
		shardWidth = DefaultShardWidth
	}

	eg := egpool.Group{PoolSize: 20}
	ids := make([]uint64, 0, len(b.ids))
	for findex, rowIDs := range b.rowIDs {
		field := b.header[findex]
		if field.Opts().Type() != FieldTypeMutex {
			continue
		}
		ids = ids[:0]

		// get slice of column ids for non-nil rowIDs and cut nil row
		// IDs out of rowIDs.
		idsIndex := 0
		for i, id := range b.ids {
			rowID := rowIDs[i]
			if rowID == nilSentinel {
				continue
			}
			rowIDs[idsIndex] = rowID
			ids = append(ids, id)
			idsIndex++
		}
		rowIDs = rowIDs[:idsIndex]

		if len(ids) == 0 {
			continue
		}
		curShard := ids[0] / shardWidth
		startIdx := 0
		for i := 1; i <= len(ids); i++ {
			var recordID uint64
			if i < len(ids) {
				recordID = ids[i]
			} else {
				recordID = (curShard + 2) * shardWidth
			}

			if recordID/shardWidth != curShard {
				endIdx := i
				shard := curShard
				field := field
				path, data, err := b.client.EncodeImport(field, shard, rowIDs[startIdx:endIdx], ids[startIdx:endIdx], false)
				if err != nil {
					return errors.Wrap(err, "encoding mutex import")
				}
				eg.Go(func() error {
					start := time.Now()
					err := b.client.DoImport(b.index.Name(), shard, path, data)
					b.log.Debugf("imp-mux %s,shard:%d,data:%d %v", field.Name(), shard, len(data), time.Since(start))
					return errors.Wrapf(err, "importing values for %s", field)
				})
				startIdx = i
				curShard = recordID / shardWidth
			}
		}
	}
	err := eg.Wait()
	if err != nil {
		if pferr := anyCause(ErrPreconditionFailed, eg.Errors()...); pferr != nil {
			return pferr
		}
		return err
	}
	return errors.Wrap(err, "importing mutex data")
}

// reset is called at the end of importing to ready the batch for the
// next round. Where possible it does not re-allocate memory.
func (b *Batch) reset() {
	b.ids = b.ids[:0]
	b.times = b.times[:0]
	for i, rowIDs := range b.rowIDs {
		b.rowIDs[i] = rowIDs[:0]
	}
	for _, tt := range b.toTranslate {
		for k := range tt {
			delete(tt, k) // TODO pool these slices
		}
	}
	for _, tts := range b.toTranslateSets {
		for k := range tts {
			delete(tts, k)
		}
	}
	for field, rowIDSet := range b.rowIDSets {
		for i := range rowIDSet {
			rowIDSet[i] = nil
		}
		b.rowIDSets[field] = rowIDSet[:0]
	}
	for _, rowIDs := range b.clearRowIDs {
		for k := range rowIDs {
			delete(rowIDs, k)
		}
	}
	for _, clearMap := range b.toTranslateClear {
		for k := range clearMap {
			delete(clearMap, k)
		}
	}
	for i := range b.toTranslateID {
		b.toTranslateID[i] = ""
	}
	for k := range b.values {
		delete(b.values, k) // TODO pool these slices
	}
	for k := range b.nullIndices {
		delete(b.nullIndices, k) // TODO pool these slices
	}
	b.cycle++
	for k, trans := range b.colTranslations {
		if trans.lastUsed-b.cycle > b.maxAge {
			delete(b.colTranslations, k)
		}
	}
	for field, rowTranslations := range b.rowTranslations {
		for k, trans := range rowTranslations {
			if trans.lastUsed-b.cycle > b.maxAge {
				delete(rowTranslations, k)
			}
		}

		if len(rowTranslations) == 0 {
			delete(b.rowTranslations, field)
		}
	}
}

// map[shard][field][view]fragmentData
type fragments map[fragmentKey]map[string]*roaring.Bitmap

type fragmentKey struct {
	shard uint64
	field string
}

func (f fragments) GetOrCreate(shard uint64, field, view string) *roaring.Bitmap {
	key := fragmentKey{shard, field}
	viewMap, ok := f[key]
	if !ok {
		viewMap = make(map[string]*roaring.Bitmap)
		f[key] = viewMap
	}
	bm, ok := viewMap[view]
	if !ok {
		bm = roaring.NewBTreeBitmap()
		viewMap[view] = bm
	}
	return bm
}

func (f fragments) GetViewMap(shard uint64, field string) map[string]*roaring.Bitmap {
	key := fragmentKey{shard, field}
	viewMap, ok := f[key]
	if !ok {
		return nil
	}
	// Remove any views which have an empty bitmap.
	// TODO: Ideally we would prevent allocating the empty bitmap to begin with,
	// but the logic is a bit tricky, and since we don't want to spend too much
	// time on it right now, we're leaving that for a future exercise.
	for k, v := range viewMap {
		if v.Count() == 0 {
			delete(viewMap, k)
		}
	}
	return viewMap
}