// Copyright 2018 Google Inc. All Rights Reserved.
//
// 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 cloudspanner

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"math/rand"
	"sort"
	"sync"
	"time"

	"cloud.google.com/go/spanner"
	"github.com/golang/glog"
	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	"github.com/google/trillian"
	"github.com/google/trillian/merkle/hashers"
	"github.com/google/trillian/storage"
	"github.com/google/trillian/storage/cache"
	"github.com/google/trillian/storage/cloudspanner/spannerpb"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

const (
	leafDataTbl            = "LeafData"
	seqDataByMerkleHashIdx = "SequenceByMerkleHash"
	seqDataTbl             = "SequencedLeafData"
	unseqTable             = "Unsequenced"

	numByteValues = 256

	unsequencedCountSQL = "SELECT Unsequenced.TreeID, COUNT(1) FROM Unsequenced GROUP BY TreeID"

	// t.TreeType: 1 = Log, 3 = PreorderedLog.
	// t.TreeState: 1 = Active, 5 = Draining.
	getActiveLogIDsSQL = `SELECT t.TreeID FROM TreeRoots t
													WHERE (t.TreeType = 1 OR t.TreeType = 3)
													AND (t.TreeState = 1 OR t.TreeState = 5)
													AND t.Deleted=false`
)

// LogStorageOptions are tuning, experiments and workarounds that can be used.
type LogStorageOptions struct {
	TreeStorageOptions

	// DequeueAcrossMerkleBuckets controls whether DequeueLeaves will only dequeue
	// from within the chosen Time+Merkle bucket, or whether it will attempt to
	// continue reading from contiguous Merkle buckets until a sufficient number
	// of leaves have been dequeued, or the entire Time bucket has been read.
	DequeueAcrossMerkleBuckets bool
	// DequeueAcrossMerkleBucketsRangeFraction specifies the fraction of Merkle
	// keyspace to dequeue from when using multi-bucket-dequeue.
	DequeueAcrossMerkleBucketsRangeFraction float64
}

var (
	// MaxUnsequencedCountStaleness configures the read-staleness limit for the
	// spanner query to retrieve the number of unsequenced certs.
	MaxUnsequencedCountStaleness = 5 * time.Minute

	// Spanner DB columns:
	colExtraData               = "ExtraData"
	colLeafValue               = "LeafValue"
	colLeafIdentityHash        = "LeafIdentityHash"
	colMerkleLeafHash          = "MerkleLeafHash"
	colSequenceNumber          = "SequenceNumber"
	colQueueTimestampNanos     = "QueueTimestampNanos"
	colIntegrateTimestampNanos = "IntegrateTimestampNanos"
)

// NewLogStorage initialises and returns a new LogStorage.
func NewLogStorage(client *spanner.Client) storage.LogStorage {
	return NewLogStorageWithOpts(client, LogStorageOptions{})
}

// NewLogStorageWithOpts initialises and returns a new LogStorage.
// The opts parameter can be used to enable custom workarounds.
func NewLogStorageWithOpts(client *spanner.Client, opts LogStorageOptions) storage.LogStorage {
	if opts.DequeueAcrossMerkleBucketsRangeFraction <= 0 || opts.DequeueAcrossMerkleBucketsRangeFraction > 1.0 {
		opts.DequeueAcrossMerkleBucketsRangeFraction = 1.0
	}
	ret := &logStorage{
		ts:   newTreeStorageWithOpts(client, opts.TreeStorageOptions),
		opts: opts,
	}

	return ret
}

// logStorage provides a Cloud Spanner backed trillian.LogStorage implementation.
// See third_party/golang/trillian/storage/log_storage.go for more details.
type logStorage struct {
	// ts provides the merkle-tree level primitives which are built upon by this
	// logStorage.
	ts *treeStorage

	// Additional options applied to this logStorage
	opts LogStorageOptions
}

func (ls *logStorage) CheckDatabaseAccessible(ctx context.Context) error {
	return checkDatabaseAccessible(ctx, ls.ts.client)
}

func (ls *logStorage) Snapshot(ctx context.Context) (storage.ReadOnlyLogTX, error) {
	var staleness spanner.TimestampBound
	if ls.opts.ReadOnlyStaleness > 0 {
		staleness = spanner.ExactStaleness(ls.opts.ReadOnlyStaleness)
	} else {
		staleness = spanner.StrongRead()
	}

	snapshotTX := &snapshotTX{
		client: ls.ts.client,
		stx:    ls.ts.client.ReadOnlyTransaction().WithTimestampBound(staleness),
		ls:     ls,
	}
	return &readOnlyLogTX{snapshotTX}, nil
}

func newLogCache(tree *trillian.Tree) (cache.SubtreeCache, error) {
	hasher, err := hashers.NewLogHasher(tree.HashStrategy)
	if err != nil {
		return cache.SubtreeCache{}, err
	}
	return cache.NewLogSubtreeCache(defLogStrata, hasher), nil
}

func (ls *logStorage) begin(ctx context.Context, tree *trillian.Tree, readonly bool, stx spanRead) (*logTX, error) {
	tx, err := ls.ts.begin(ctx, tree, newLogCache, stx)
	if err != nil {
		return nil, err
	}

	// Sanity check tx.config
	if cfg, ok := tx.config.(*spannerpb.LogStorageConfig); !ok || cfg == nil {
		return nil, fmt.Errorf("unexpected config type for LOG tree %v: %T", tx.treeID, tx.config)
	}

	return &logTX{
		ls:       ls,
		dequeued: make(map[string]*QueuedEntry),
		treeTX:   tx,
	}, nil
}

func (ls *logStorage) BeginForTree(ctx context.Context, treeID int64) (storage.LogTreeTX, error) {
	return nil, ErrNotImplemented
}

func (ls *logStorage) ReadWriteTransaction(ctx context.Context, tree *trillian.Tree, f storage.LogTXFunc) error {
	_, err := ls.ts.client.ReadWriteTransaction(ctx, func(ctx context.Context, stx *spanner.ReadWriteTransaction) error {
		tx, err := ls.begin(ctx, tree, false /* readonly */, stx)
		if err != nil {
			return err
		}
		if err := f(ctx, tx); err != nil {
			return err
		}
		return tx.flushSubtrees()
	})
	return err
}

func (ls *logStorage) SnapshotForTree(ctx context.Context, tree *trillian.Tree) (storage.ReadOnlyLogTreeTX, error) {
	return ls.begin(ctx, tree, true /* readonly */, ls.ts.client.ReadOnlyTransaction())
}

func (ls *logStorage) QueueLeaves(ctx context.Context, tree *trillian.Tree, leaves []*trillian.LogLeaf, qTimestamp time.Time) ([]*trillian.QueuedLogLeaf, error) {
	_, treeConfig, err := ls.ts.getTreeAndConfig(ctx, tree)
	if err != nil {
		return nil, err
	}
	config, ok := treeConfig.(*spannerpb.LogStorageConfig)
	if !ok {
		return nil, status.Errorf(codes.Internal, "got unexpected config type for Log operation: %T", treeConfig)
	}

	now := time.Now().UTC().Unix()
	bucketPrefix := (now % config.NumUnseqBuckets) << 8

	results := make([]*trillian.QueuedLogLeaf, len(leaves))
	writeDupes := make(map[string][]indexMerkleHash)

	qTS := qTimestamp.UnixNano()
	var wg sync.WaitGroup
	for i, l := range leaves {
		wg.Add(1)
		// Capture values of i and l for later reference in the MutationResultFunc below.
		i := i
		l := l
		go func() {
			defer wg.Done()

			// The insert of the leafdata and the unsequenced work item must happen
			// atomically.
			m1 := spanner.Insert(
				leafDataTbl,
				[]string{colTreeID, colLeafIdentityHash, colLeafValue, colExtraData, colQueueTimestampNanos},
				[]interface{}{tree.TreeId, l.LeafIdentityHash, l.LeafValue, l.ExtraData, qTS})
			b := bucketPrefix | int64(l.MerkleLeafHash[0])
			m2 := spanner.Insert(
				unseqTable,
				[]string{colTreeID, colBucket, colQueueTimestampNanos, colMerkleLeafHash, colLeafIdentityHash},
				[]interface{}{tree.TreeId, b, qTS, l.MerkleLeafHash, l.LeafIdentityHash})

			_, err = ls.ts.client.Apply(ctx, []*spanner.Mutation{m1, m2})
			if spanner.ErrCode(err) == codes.AlreadyExists {
				k := string(l.LeafIdentityHash)
				writeDupes[k] = append(writeDupes[k], indexMerkleHash{i, l.MerkleLeafHash})
			} else if err != nil {
				s, _ := status.FromError(err)
				results[i] = &trillian.QueuedLogLeaf{Status: s.Proto()}
			}
		}()
	}

	// Wait for all of our mutations to apply (or fail):
	wg.Wait()

	// Finally, read back any leaves which failed with an already exists error
	// when we tried to insert them:
	err = ls.readDupeLeaves(ctx, tree.TreeId, writeDupes, results)
	if err != nil {
		return nil, err
	}
	return results, nil
}

func (ls *logStorage) AddSequencedLeaves(ctx context.Context, tree *trillian.Tree, leaves []*trillian.LogLeaf) ([]*trillian.QueuedLogLeaf, error) {
	return nil, ErrNotImplemented
}

// readDupeLeaves reads the leaves whose ids are passed as keys in the dupes map,
// and stores them in results.
func (ls *logStorage) readDupeLeaves(ctx context.Context, logID int64, dupes map[string][]indexMerkleHash, results []*trillian.QueuedLogLeaf) error {
	numDupes := len(dupes)
	if numDupes == 0 {
		return nil
	}
	glog.V(2).Infof("dupe rowsToRead: %v", numDupes)

	ids := make([][]byte, 0, numDupes)
	for k := range dupes {
		ids = append(ids, []byte(k))
	}
	dupesRead := 0
	tx := ls.ts.client.Single()
	err := readLeaves(ctx, tx, logID, ids, func(l *trillian.LogLeaf) {
		glog.V(2).Infof("Found already exists dupe: %v", l)
		dupesRead++

		indices := dupes[string(l.LeafIdentityHash)]
		glog.V(2).Infof("Indices %v", indices)
		if len(indices) == 0 {
			glog.Warningf("Logic error: Spanner returned a leaf %x, but it matched no requested index", l.LeafIdentityHash)
			return
		}
		for _, i := range indices {
			leaf := l
			leaf.MerkleLeafHash = i.merkleHash
			results[i.index] = &trillian.QueuedLogLeaf{
				Leaf:   leaf,
				Status: status.Newf(codes.AlreadyExists, "leaf already exists: %v", l.LeafIdentityHash).Proto(),
			}
		}
	})
	tx.Close()
	if err != nil {
		return err
	}
	if got, want := dupesRead, numDupes; got != want {
		return fmt.Errorf("read unexpected number of dupe rows %d, want %d", got, want)
	}
	return nil
}

// logTX is a concrete implementation of the Trillian storage.LogStorage
// interface.
type logTX struct {
	// treeTX embeds the merkle-tree level transactional actions.
	*treeTX

	// logStorage is the logStorage which begat this logTX.
	ls *logStorage

	// numSequenced holds the number of leaves sequenced by this transaction.
	numSequenced int64

	// dequeued is a map of LeafIdentityHash to QueuedEntry containing entries for
	// everything dequeued by this transaction.
	// This is required to recover the primary key for the unsequenced entry in
	// UpdateSequencedLeaves.
	dequeued map[string]*QueuedEntry
}

func (tx *logTX) getLogStorageConfig() *spannerpb.LogStorageConfig {
	return tx.config.(*spannerpb.LogStorageConfig)
}

// LatestSignedLogRoot returns the freshest SignedLogRoot for this log at the
// time the transaction was started.
func (tx *logTX) LatestSignedLogRoot(ctx context.Context) (trillian.SignedLogRoot, error) {
	currentSTH, err := tx.currentSTH(ctx)
	if err != nil {
		return trillian.SignedLogRoot{}, err
	}
	writeRev, err := tx.writeRev(ctx)
	if err != nil {
		return trillian.SignedLogRoot{}, err
	}

	if got, want := currentSTH.TreeRevision+1, writeRev; got != want {
		return trillian.SignedLogRoot{}, fmt.Errorf("inconsistency: currentSTH.TreeRevision+1 (%d) != writeRev (%d)", got, want)
	}

	apiSig, err := storageToAPISig(currentSTH.Signature)
	if err != nil {
		return trillian.SignedLogRoot{}, err
	}
	// We already read the latest root as part of starting the transaction (in
	// order to calculate the writeRevision), so we just return that data here:
	return trillian.SignedLogRoot{
		TimestampNanos: currentSTH.TsNanos,
		RootHash:       currentSTH.RootHash,
		TreeSize:       currentSTH.TreeSize,
		LogId:          currentSTH.TreeId,
		TreeRevision:   currentSTH.TreeRevision,
		Signature:      apiSig,
	}, nil
}

// StoreSignedLogRoot stores the provided root.
// This method will return an error if the caller attempts to store more than
// one root per log for a given tree size.
func (tx *logTX) StoreSignedLogRoot(ctx context.Context, root trillian.SignedLogRoot) error {
	writeRev, err := tx.writeRev(ctx)
	if err == storage.ErrTreeNeedsInit {
		writeRev = 0
	} else if err != nil {
		return err
	}

	if got, want := root.LogId, tx.treeID; got != want {
		return fmt.Errorf("provided root is for LogID (%d) != transaction's treeID (%d)", got, want)
	}
	storageSig, err := apiToStorageSig(root.Signature)
	if err != nil {
		return err
	}
	sth := spannerpb.TreeHead{
		TsNanos:      root.TimestampNanos,
		RootHash:     root.RootHash,
		TreeSize:     root.TreeSize,
		TreeId:       tx.treeID,
		TreeRevision: writeRev,
		Signature:    storageSig,
	}
	var sigBytes, metaBytes []byte
	if sth.Signature == nil {
		return errors.New("sth signature is nil")
	}
	sigBytes, err = proto.Marshal(sth.Signature)
	if err != nil {
		return err
	}
	if sth.Metadata != nil {
		metaBytes, err = proto.Marshal(sth.Metadata)
		if err != nil {
			return err
		}
	}

	m := spanner.Insert(
		"TreeHeads",
		[]string{
			"TreeID",
			"TimestampNanos",
			"TreeSize",
			"RootHash",
			"RootSignature",
			"TreeRevision",
			"TreeMetadata",
		},
		[]interface{}{
			sth.TreeId,
			sth.TsNanos,
			sth.TreeSize,
			sth.RootHash,
			sigBytes,
			sth.TreeRevision,
			metaBytes,
		})

	stx, ok := tx.stx.(*spanner.ReadWriteTransaction)
	if !ok {
		return ErrWrongTXType
	}
	return stx.BufferWrite([]*spanner.Mutation{m})
}

func readLeaves(ctx context.Context, stx *spanner.ReadOnlyTransaction, logID int64, ids [][]byte, f func(*trillian.LogLeaf)) error {
	leafTable := leafDataTbl
	cols := []string{colLeafIdentityHash, colLeafValue, colExtraData, colQueueTimestampNanos}
	keys := make([]spanner.KeySet, 0)
	for _, l := range ids {
		keys = append(keys, spanner.Key{logID, l})
	}

	rows := stx.Read(ctx, leafTable, spanner.KeySets(keys...), cols)
	return rows.Do(func(r *spanner.Row) error {
		var l trillian.LogLeaf
		var qTimestamp int64
		if err := r.Columns(&l.LeafIdentityHash, &l.LeafValue, &l.ExtraData, &qTimestamp); err != nil {
			return err
		}
		var err error
		l.QueueTimestamp, err = ptypes.TimestampProto(time.Unix(0, qTimestamp))
		if err != nil {
			return fmt.Errorf("got invalid queue timestamp: %v", err)
		}
		f(&l)
		return nil
	})
}

type indexMerkleHash struct {
	index      int
	merkleHash []byte
}

func (tx *logTX) QueueLeaves(ctx context.Context, leaves []*trillian.LogLeaf, ts time.Time) ([]*trillian.LogLeaf, error) {
	return nil, ErrNotImplemented
}

// genDequeueBucketsSingleSplit returns a number of bucket IDs suitable for using in
// DequeueLeaves to determine the list of bucket IDs to sequence from.
//
// The IDs are 16-bits long, the high byte is an offset into the time ring and will
// be on the opposite side of the ring to where the value in now would normally
// lie, and the low byte is a merkle hash prefix.
func genDequeueBucketsSingleSplit(config *spannerpb.LogStorageConfig, now int64, merkleBucket int32) []int32 {
	n := int32(numByteValues / config.NumMerkleBuckets)
	ret := make([]int32, 0, n)
	bucketHigh := int32((((now + config.NumUnseqBuckets/2) % config.NumUnseqBuckets) << 8))
	for i := merkleBucket; i < merkleBucket+n; i++ {
		ret = append(ret, bucketHigh|(i%numByteValues))
	}
	return ret
}

// genDequeueRowsOld returns a slice containing a single RowSet of bucket prefixes to dequeue from.
func genDequeueRowsOld(config *spannerpb.LogStorageConfig, treeID int64, now int64, merkleBucket int32) []spanner.KeySet {
	rows := make([]spanner.KeySet, 0)
	for _, b := range genDequeueBucketsSingleSplit(config, now, merkleBucket) {
		rows = append(rows, spanner.Key{treeID, int64(b)}.AsPrefix())
	}
	return []spanner.KeySet{spanner.KeySets(rows...)}
}

// genDequeueAcrossBuckets returns a slice containing two KeySets to dequeue from.
//
// This function returns a pair of KeySets which cover the entirety of the time bucket
// implied by now.  The first of these is from [time,merkleBucket] to the end of the
// time bucket range, and the second from the beginning of the time bucket to [time, merkleBucket).
//
// In general, trying to sequence many leaves across all of these Merkle buckets would
// be a bad idea, because the Merkle buckets are across several splits in the DB, however,
// where there are few unsequenced leaves it may be beneficial to take the hit on the
// coordinated commit since it should be faster than having the sequencer iterate many
// times, doing very little work each time, to pick out all the individual entries.
//
// The returned KeySets should be read *in order* only, and the read call should be
// bailed out once a sufficient number of leaves have been read.  This should mean that
// when there are many leaves in the queue, the behaviour tends back towards the original
// behaviour of picking leaves from a very small number of Merkle buckets.
func genDequeueRowsAcrossBuckets(config *spannerpb.LogStorageConfig, treeID int64, now int64, merkleBucket int32, bucketFraction float64) []spanner.KeySet {
	merkleBucket = merkleBucket % numByteValues
	merkleBucketLimit := merkleBucket + int32(numByteValues*bucketFraction)

	bucketHigh := int32((((now + config.NumUnseqBuckets/2) % config.NumUnseqBuckets) << 8))

	if merkleBucketLimit <= 0xff {
		return []spanner.KeySet{spanner.KeyRange{Kind: spanner.ClosedClosed, Start: spanner.Key{treeID, int64(bucketHigh | merkleBucket)}, End: spanner.Key{treeID, int64(bucketHigh | merkleBucketLimit)}}}
	}

	// The range is too big and wraps around, overflowing a byte value, so we'll
	// start the second range at 0x00 and end at the upper limit modulo 256:
	merkleBucketLimit %= 0x100
	return []spanner.KeySet{
		spanner.KeyRange{Kind: spanner.ClosedClosed, Start: spanner.Key{treeID, bucketHigh | merkleBucket}, End: spanner.Key{treeID, bucketHigh | 0xff}},
		spanner.KeyRange{Kind: spanner.ClosedClosed, Start: spanner.Key{treeID, bucketHigh | 0x00}, End: spanner.Key{treeID, bucketHigh | merkleBucketLimit}},
	}

}

// DequeueLeaves removes [0, limit) leaves from the to-be-sequenced queue.
// The leaves returned are not guaranteed to be in any particular order.
// The caller should assign sequence numbers and pass the updated leaves as
// arguments to the UpdateSequencedLeaves method.
//
// The LogLeaf structs returned by this method will not be fully populated;
// only the LeafIdentityHash and MerkleLeafHash fields will contain data, this
// should be sufficient for assigning sequence numbers with this storage impl.
//
// TODO(al): cutoff is currently ignored.
func (tx *logTX) DequeueLeaves(ctx context.Context, limit int, cutoff time.Time) ([]*trillian.LogLeaf, error) {
	if limit <= 0 {
		return nil, fmt.Errorf("limit should be > 0, got %d", limit)
	}

	// Decide which bucket(s) to dequeue from.
	// The first part of the bucket key is a time based ring - at any given
	// moment, FEs queueing entries will be adding them to different buckets
	// than we're dequeuing from here.
	now := time.Now().UTC()
	var rows []spanner.KeySet

	if tx.ls.opts.DequeueAcrossMerkleBuckets {
		glog.V(1).Info("Using multi-bucket dequeue")
		rows = genDequeueRowsAcrossBuckets(tx.getLogStorageConfig(), tx.treeID, now.Unix(), rand.Int31n(numByteValues), tx.ls.opts.DequeueAcrossMerkleBucketsRangeFraction)
	} else {
		rows = genDequeueRowsOld(tx.getLogStorageConfig(), tx.treeID, now.Unix(), rand.Int31n(numByteValues))
	}
	ret := make([]*trillian.LogLeaf, 0, limit)
	for _, rs := range rows {
		rows := tx.stx.Read(ctx, unseqTable, rs, []string{colBucket, colQueueTimestampNanos, colMerkleLeafHash, colLeafIdentityHash})
		if err := rows.Do(func(r *spanner.Row) error {
			if len(ret) >= limit {
				return errFinished
			}
			var l trillian.LogLeaf
			var qe QueuedEntry
			if err := r.Columns(&qe.bucket, &qe.timestamp, &l.MerkleLeafHash, &l.LeafIdentityHash); err != nil {
				return err
			}

			var err error
			l.QueueTimestamp, err = ptypes.TimestampProto(time.Unix(0, qe.timestamp))
			if err != nil {
				return fmt.Errorf("got invalid queue timestamp: %v", err)
			}
			k := string(l.LeafIdentityHash)
			if tx.dequeued[k] != nil {
				// dupe, user probably called DequeueLeaves more than once.
				return nil
			}

			ret = append(ret, &l)
			qe.leaf = &l
			tx.dequeued[k] = &qe
			return nil
		}); err != nil && err != errFinished {
			return nil, err
		}

		// If we've already got enough leaves, don't wrap around for any further reads.
		if len(ret) >= limit {
			break
		}
	}
	return ret, nil
}

// UpdateSequencedLeaves stores the sequence numbers assigned to the leaves,
// and integrates them into the tree.
func (tx *logTX) UpdateSequencedLeaves(ctx context.Context, leaves []*trillian.LogLeaf) error {
	stx, ok := tx.stx.(*spanner.ReadWriteTransaction)
	if !ok {
		return ErrWrongTXType
	}
	// We need the latest root to know what the next sequence number to use below is.
	currentSTH, err := tx.currentSTH(ctx)
	if err != nil {
		return err
	}

	for _, l := range leaves {
		if got, want := l.LeafIndex, currentSTH.TreeSize+tx.numSequenced; got != want {
			return fmt.Errorf("attempting to assign non-sequential leaf with sequence %d, want %d", got, want)
		}

		qe, ok := tx.dequeued[string(l.LeafIdentityHash)]
		if !ok {
			return fmt.Errorf("attempting to assign unknown merkleleafhash %v", l.MerkleLeafHash)
		}

		iTimestamp, err := ptypes.Timestamp(l.IntegrateTimestamp)
		if err != nil {
			return fmt.Errorf("got invalid integrate timestamp: %v", err)
		}

		// Add the sequence mapping...
		m1 := spanner.Insert(seqDataTbl,
			[]string{colTreeID, colSequenceNumber, colLeafIdentityHash, colMerkleLeafHash, colIntegrateTimestampNanos},
			[]interface{}{tx.treeID, l.LeafIndex, l.LeafIdentityHash, l.MerkleLeafHash, iTimestamp.UnixNano()})

		m2 := spanner.Delete(unseqTable, spanner.Key{tx.treeID, qe.bucket, qe.timestamp, l.MerkleLeafHash})

		tx.numSequenced++
		if err := stx.BufferWrite([]*spanner.Mutation{m1, m2}); err != nil {
			return fmt.Errorf("bufferwrite(): %v", err)
		}
	}

	return nil
}

// GetSequencedLeafCount returns the number of leaves integrated into the tree
// at the time the transaction was started.
func (tx *logTX) GetSequencedLeafCount(ctx context.Context) (int64, error) {
	currentSTH, err := tx.currentSTH(ctx)
	if err != nil {
		return -1, err
	}

	return currentSTH.TreeSize, nil
}

// leafmap is a map of LogLeaf by sequence number which knows how to populate
// itself directly from Spanner Rows.
type leafmap map[int64]*trillian.LogLeaf

// addFullRow appends the leaf data in row to the array
func (l leafmap) addFullRow(r *spanner.Row) error {
	var (
		merkleLeafHash, leafValue, extraData []byte
		sequenceNumber                       int64
		leafIDHash                           []byte
		qTimestamp                           int64
		iTimestamp                           int64
	)

	//`SELECT sd.MerkleLeafHash, ld.LeafValue, ld.ExtraData, sd.SequenceNumber, ld.LeafIdentityHash, ld.QueueTimestampNanos, sd.IntegrateTimestampNanos
	if err := r.Columns(&merkleLeafHash, &leafValue, &extraData, &sequenceNumber, &leafIDHash, &qTimestamp, &iTimestamp); err != nil {
		return err
	}
	leaf := &trillian.LogLeaf{
		MerkleLeafHash:   merkleLeafHash,
		LeafValue:        leafValue,
		ExtraData:        extraData,
		LeafIndex:        sequenceNumber,
		LeafIdentityHash: leafIDHash,
	}
	var err error
	leaf.QueueTimestamp, err = ptypes.TimestampProto(time.Unix(0, qTimestamp))
	if err != nil {
		return fmt.Errorf("got invalid queue timestamp %v", err)
	}
	leaf.IntegrateTimestamp, err = ptypes.TimestampProto(time.Unix(0, iTimestamp))
	if err != nil {
		return fmt.Errorf("got invalid integrate timestamp %v", err)
	}

	l[sequenceNumber] = leaf
	return nil
}

// leavesByHash is a map of []LogLeaf (keyed by value hash) which knows how to
// populate itself from Spanner Rows.
type leavesByHash map[string][]*trillian.LogLeaf

// addRow adds the contents of the Spanner Row to this map.
func (b leavesByHash) addRow(r *spanner.Row) error {
	var h []byte
	var v []byte
	var ed []byte
	var qTimestamp int64
	if err := r.Columns(&h, &v, &ed, &qTimestamp); err != nil {
		return err
	}
	queueTimestamp, err := ptypes.TimestampProto(time.Unix(0, qTimestamp))
	if err != nil {
		return fmt.Errorf("got invalid queue timestamp: %v", err)
	}

	leaves, ok := b[string(h)]
	if !ok {
		return fmt.Errorf("inconsistency: unexpected leafValueHash %v", h)
	}
	for i := range leaves {
		if got, want := leaves[i].LeafIdentityHash, h; !bytes.Equal(got, want) {
			return fmt.Errorf("inconsistency: unexpected leafvaluehash %v, want %v", got, want)
		}
		leaves[i].LeafValue = v
		leaves[i].ExtraData = ed
		leaves[i].QueueTimestamp = queueTimestamp
	}
	return nil
}

// populateLeafData populates the partial LogLeaf structs held in the passed in
// map of LeafIdentityHash to []LogLeaf by reading the remaining LogLeaf data from
// Spanner.
// The value of byHash is an []LogLeaf because the underlying leaf data could
// be sequenced into multiple tree leaves if the log allows duplication.
func (tx *logTX) populateLeafData(ctx context.Context, byHash leavesByHash) error {
	keySet := make([]spanner.KeySet, 0, len(byHash))
	for k := range byHash {
		keySet = append(keySet, spanner.Key{tx.treeID, []byte(k)})
	}
	cols := []string{colLeafIdentityHash, colLeafValue, colExtraData, colQueueTimestampNanos}
	rows := tx.stx.Read(ctx, leafDataTbl, spanner.KeySets(keySet...), cols)
	return rows.Do(byHash.addRow)
}

// validateIndices ensures that all indices are between 0 and treeSize-1.
func validateIndices(indices []int64, treeSize int64) error {
	maxIndex := treeSize - 1
	for _, i := range indices {
		if i < 0 {
			return status.Errorf(codes.InvalidArgument, "index %d is < 0", i)
		}
		if i > maxIndex {
			return status.Errorf(codes.OutOfRange, "index %d is > highest index in current tree %d", i, maxIndex)
		}
	}
	return nil
}

// GetLeavesByIndex returns the leaves corresponding to the given indices.
func (tx *logTX) GetLeavesByIndex(ctx context.Context, indices []int64) ([]*trillian.LogLeaf, error) {
	// We need the latest root to validate the indices are within range.
	currentSTH, err := tx.currentSTH(ctx)
	if err != nil {
		return nil, err
	}

	if err := validateIndices(indices, currentSTH.TreeSize); err != nil {
		return nil, err
	}

	leaves := make(leafmap)
	stmt := spanner.NewStatement(
		`SELECT sd.MerkleLeafHash, ld.LeafValue, ld.ExtraData, sd.SequenceNumber, ld.LeafIdentityHash, ld.QueueTimestampNanos, sd.IntegrateTimestampNanos
FROM SequencedLeafData as sd
INNER JOIN LeafData as ld
ON sd.TreeID = ld.TreeID AND sd.LeafIdentityHash = ld.LeafIdentityHash
WHERE sd.TreeID = @tree_id and sd.SequenceNumber IN UNNEST(@seq_nums)`)
	stmt.Params["tree_id"] = tx.treeID
	stmt.Params["seq_nums"] = indices

	rows := tx.stx.Query(ctx, stmt)
	if err := rows.Do(leaves.addFullRow); err != nil {
		return nil, err
	}

	// Sanity check that we got everything we wanted
	if got, want := len(leaves), len(indices); got != want {
		return nil, fmt.Errorf("inconsistency: got %d leaves, want %d", got, want)
	}

	// Sort the leaves so they are in the same order as the indices.
	ret := make([]*trillian.LogLeaf, 0, len(indices))
	for _, i := range indices {
		l, ok := leaves[i]
		if !ok {
			return nil, fmt.Errorf("inconsistency: missing data for index %d", i)
		}
		ret = append(ret, l)
	}

	return ret, nil
}

func validateRange(start, count, treeSize int64) error {
	if count <= 0 {
		return status.Errorf(codes.InvalidArgument, "invalid count %d", count)
	}
	if start < 0 {
		return status.Errorf(codes.InvalidArgument, "invalid start %d", start)
	}
	if start >= treeSize {
		return status.Errorf(codes.OutOfRange, "start index %d beyond tree size %d", start, treeSize)
	}
	return nil
}

// GetLeavesByRange returns the leaves corresponding to the given index range.
func (tx *logTX) GetLeavesByRange(ctx context.Context, start, count int64) ([]*trillian.LogLeaf, error) {
	// We need the latest root to validate the indices are within range.
	currentSTH, err := tx.currentSTH(ctx)
	if err != nil {
		return nil, err
	}

	if err := validateRange(start, count, currentSTH.TreeSize); err != nil {
		return nil, err
	}

	stmt := spanner.NewStatement(
		`SELECT sd.MerkleLeafHash, ld.LeafValue, ld.ExtraData, sd.SequenceNumber, ld.LeafIdentityHash, ld.QueueTimestampNanos, sd.IntegrateTimestampNanos
FROM SequencedLeafData as sd
INNER JOIN LeafData as ld
ON sd.TreeID = ld.TreeID AND sd.LeafIdentityHash = ld.LeafIdentityHash
WHERE sd.TreeID = @tree_id AND sd.SequenceNumber >= @start AND sd.SequenceNumber < @xend`)
	stmt.Params["tree_id"] = tx.treeID
	stmt.Params["start"] = start
	xend := start + count
	if xend > currentSTH.TreeSize {
		xend = currentSTH.TreeSize
		count = xend - start
	}
	stmt.Params["xend"] = xend

	// Results need to be returned in order [start, end), all of which should be
	// available (as we restricted xend/count to TreeSize).
	leaves := make(leafmap)
	rows := tx.stx.Query(ctx, stmt)
	if err := rows.Do(leaves.addFullRow); err != nil {
		return nil, err
	}
	ret := make([]*trillian.LogLeaf, 0, count)
	for i := start; i < (start + count); i++ {
		l, ok := leaves[i]
		if !ok {
			return nil, fmt.Errorf("inconsistency: missing data for index %d", i)
		}
		ret = append(ret, l)
		delete(leaves, i)
	}
	if len(leaves) > 0 {
		return nil, fmt.Errorf("inconsistency: unexpected extra data outside range %d, +%d", start, count)
	}

	return ret, nil
}

// leafSlice is a slice of LogLeaf which knows how to populate itself from
// Spanner Rows.
type leafSlice []*trillian.LogLeaf

// addRow appends the leaf data in Row to the array.
func (l *leafSlice) addRow(r *spanner.Row) error {
	var (
		s      int64
		mh, lh []byte
	)

	if err := r.Columns(&s, &mh, &lh); err != nil {
		return err
	}
	leaf := trillian.LogLeaf{
		LeafIndex:        s,
		MerkleLeafHash:   mh,
		LeafIdentityHash: lh,
	}
	*l = append(*l, &leaf)
	return nil
}

// getUsingIndex returns a slice containing the LogLeaf structs corresponding
// to the requested keys.
// The entries in key are used in constructing a primary key (treeID, keyElem)
// for the specified Spanner index.
// If bySeq is true, the returned slice will be order by LogLeaf.LeafIndex.
func (tx *logTX) getUsingIndex(ctx context.Context, idx string, keys [][]byte, bySeq bool) ([]*trillian.LogLeaf, error) {
	keySet := make([]spanner.KeySet, 0, len(keys))
	for _, k := range keys {
		keySet = append(keySet, spanner.Key{tx.treeID, k})
	}

	leaves := make(leafSlice, 0, len(keys))
	cols := []string{colSequenceNumber, colMerkleLeafHash, colLeafIdentityHash}
	rows := tx.stx.ReadUsingIndex(ctx, seqDataTbl, idx, spanner.KeySets(keySet...), cols)
	if err := rows.Do(leaves.addRow); err != nil {
		return nil, err
	}

	byHash := make(leavesByHash)
	for i := range leaves {
		k := string(leaves[i].LeafIdentityHash)
		byHash[k] = append(byHash[k], leaves[i])
	}

	// Now we can fetch & combine the actual leaf data:
	if err := tx.populateLeafData(ctx, byHash); err != nil {
		return nil, err
	}

	if bySeq {
		sort.Sort(byIndex(leaves))
	}

	return leaves, nil
}

// GetLeavesByHash returns the leaves corresponding to the given merkle hashes.
// Any unknown hashes will simply be ignored, and the caller should inspect the
// returned leaves to determine whether this has occurred.
// TODO(al): Currently, this method does not populate the IntegrateTimestamp
//   member of the returned leaves. We should convert this method to use SQL
//   rather than denormalising IntegrateTimestampNanos into the index too.
func (tx *logTX) GetLeavesByHash(ctx context.Context, hashes [][]byte, bySeq bool) ([]*trillian.LogLeaf, error) {
	return tx.getUsingIndex(ctx, seqDataByMerkleHashIdx, hashes, bySeq)
}

// QueuedEntry represents a leaf which was dequeued.
// It's used to store some extra info which is necessary for rebuilding the
// leaf's primary key when it's passed back in to UpdateSequencedLeaves.
type QueuedEntry struct {
	// leaf is partially populated with the Merkle and LeafValue hashes only.
	leaf      *trillian.LogLeaf
	bucket    int64
	timestamp int64
}

// readOnlyLogTX implements storage.ReadOnlyLogTX.
type readOnlyLogTX struct {
	*snapshotTX
}

func (tx *readOnlyLogTX) GetActiveLogIDs(ctx context.Context) ([]int64, error) {
	tx.mu.RLock()
	defer tx.mu.RUnlock()
	if tx.stx == nil {
		return nil, ErrTransactionClosed
	}

	ids := []int64{}
	// We have to use SQL as Read() doesn't work against an index.
	stmt := spanner.NewStatement(getActiveLogIDsSQL)
	rows := tx.stx.Query(ctx, stmt)
	if err := rows.Do(func(r *spanner.Row) error {
		var id int64
		if err := r.Columns(&id); err != nil {
			return err
		}
		ids = append(ids, id)
		return nil
	}); err != nil {
		glog.Warning("GetActiveLogIDs: %v", err)
		return nil, fmt.Errorf("problem executing getActiveLogIDsSQL: %v", err)
	}
	return ids, nil
}

func (tx *readOnlyLogTX) GetUnsequencedCounts(ctx context.Context) (storage.CountByLogID, error) {
	stmt := spanner.NewStatement(unsequencedCountSQL)
	ret := make(storage.CountByLogID)
	rows := tx.stx.Query(ctx, stmt)
	if err := rows.Do(func(r *spanner.Row) error {
		var id, c int64
		if err := r.Columns(&id, &c); err != nil {
			return err
		}
		ret[id] = c
		return nil
	}); err != nil {
		return nil, fmt.Errorf("problem executing unsequencedCountSQL: %v", err)
	}
	return ret, nil
}

// LogLeaf sorting boilerplate below.

type byIndex []*trillian.LogLeaf

func (b byIndex) Len() int { return len(b) }

func (b byIndex) Swap(i, j int) { b[i], b[j] = b[j], b[i] }

func (b byIndex) Less(i, j int) bool { return b[i].LeafIndex < b[j].LeafIndex }
