/*
Copyright 2017 Google LLC

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 spanner

import (
	"bytes"
	"context"
	"io"
	"log"
	"sync/atomic"
	"time"

	"cloud.google.com/go/internal/protostruct"
	"cloud.google.com/go/internal/trace"
	"github.com/golang/protobuf/proto"
	proto3 "github.com/golang/protobuf/ptypes/struct"
	"github.com/googleapis/gax-go/v2"
	"google.golang.org/api/iterator"
	sppb "google.golang.org/genproto/googleapis/spanner/v1"
	"google.golang.org/grpc/codes"
)

// streamingReceiver is the interface for receiving data from a client side
// stream.
type streamingReceiver interface {
	Recv() (*sppb.PartialResultSet, error)
}

// errEarlyReadEnd returns error for read finishes when gRPC stream is still
// active.
func errEarlyReadEnd() error {
	return spannerErrorf(codes.FailedPrecondition, "read completed with active stream")
}

// stream is the internal fault tolerant method for streaming data from Cloud
// Spanner.
func stream(ctx context.Context, logger *log.Logger, rpc func(ct context.Context, resumeToken []byte) (streamingReceiver, error), setTimestamp func(time.Time), release func(error)) *RowIterator {
	ctx, cancel := context.WithCancel(ctx)
	ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.RowIterator")
	return &RowIterator{
		streamd:      newResumableStreamDecoder(ctx, logger, rpc),
		rowd:         &partialResultSetDecoder{},
		setTimestamp: setTimestamp,
		release:      release,
		cancel:       cancel,
	}
}

// RowIterator is an iterator over Rows.
type RowIterator struct {
	// The plan for the query. Available after RowIterator.Next returns
	// iterator.Done if QueryWithStats was called.
	QueryPlan *sppb.QueryPlan

	// Execution statistics for the query. Available after RowIterator.Next
	// returns iterator.Done if QueryWithStats was called.
	QueryStats map[string]interface{}

	// For a DML statement, the number of rows affected. For PDML, this is a
	// lower bound. Available for DML statements after RowIterator.Next returns
	// iterator.Done.
	RowCount int64

	streamd      *resumableStreamDecoder
	rowd         *partialResultSetDecoder
	setTimestamp func(time.Time)
	release      func(error)
	cancel       func()
	err          error
	rows         []*Row
	sawStats     bool
}

// Next returns the next result. Its second return value is iterator.Done if
// there are no more results. Once Next returns Done, all subsequent calls
// will return Done.
func (r *RowIterator) Next() (*Row, error) {
	if r.err != nil {
		return nil, r.err
	}
	for len(r.rows) == 0 && r.streamd.next() {
		prs := r.streamd.get()
		if prs.Stats != nil {
			r.sawStats = true
			r.QueryPlan = prs.Stats.QueryPlan
			r.QueryStats = protostruct.DecodeToMap(prs.Stats.QueryStats)
			if prs.Stats.RowCount != nil {
				rc, err := extractRowCount(prs.Stats)
				if err != nil {
					return nil, err
				}
				r.RowCount = rc
			}
		}
		r.rows, r.err = r.rowd.add(prs)
		if r.err != nil {
			return nil, r.err
		}
		if !r.rowd.ts.IsZero() && r.setTimestamp != nil {
			r.setTimestamp(r.rowd.ts)
			r.setTimestamp = nil
		}
	}
	if len(r.rows) > 0 {
		row := r.rows[0]
		r.rows = r.rows[1:]
		return row, nil
	}
	if err := r.streamd.lastErr(); err != nil {
		r.err = toSpannerError(err)
	} else if !r.rowd.done() {
		r.err = errEarlyReadEnd()
	} else {
		r.err = iterator.Done
	}
	return nil, r.err
}

func extractRowCount(stats *sppb.ResultSetStats) (int64, error) {
	if stats.RowCount == nil {
		return 0, spannerErrorf(codes.Internal, "missing RowCount")
	}
	switch rc := stats.RowCount.(type) {
	case *sppb.ResultSetStats_RowCountExact:
		return rc.RowCountExact, nil
	case *sppb.ResultSetStats_RowCountLowerBound:
		return rc.RowCountLowerBound, nil
	default:
		return 0, spannerErrorf(codes.Internal, "unknown RowCount type %T", stats.RowCount)
	}
}

// Do calls the provided function once in sequence for each row in the
// iteration. If the function returns a non-nil error, Do immediately returns
// that error.
//
// If there are no rows in the iterator, Do will return nil without calling the
// provided function.
//
// Do always calls Stop on the iterator.
func (r *RowIterator) Do(f func(r *Row) error) error {
	defer r.Stop()
	for {
		row, err := r.Next()
		switch err {
		case iterator.Done:
			return nil
		case nil:
			if err = f(row); err != nil {
				return err
			}
		default:
			return err
		}
	}
}

// Stop terminates the iteration. It should be called after you finish using the
// iterator.
func (r *RowIterator) Stop() {
	if r.streamd != nil {
		defer trace.EndSpan(r.streamd.ctx, r.err)
	}
	if r.cancel != nil {
		r.cancel()
	}
	if r.release != nil {
		r.release(r.err)
		if r.err == nil {
			r.err = spannerErrorf(codes.FailedPrecondition, "Next called after Stop")
		}
		r.release = nil

	}
}

// partialResultQueue implements a simple FIFO queue.  The zero value is a valid
// queue.
type partialResultQueue struct {
	q     []*sppb.PartialResultSet
	first int
	last  int
	n     int // number of elements in queue
}

// empty returns if the partialResultQueue is empty.
func (q *partialResultQueue) empty() bool {
	return q.n == 0
}

// errEmptyQueue returns error for dequeuing an empty queue.
func errEmptyQueue() error {
	return spannerErrorf(codes.OutOfRange, "empty partialResultQueue")
}

// peekLast returns the last item in partialResultQueue; if the queue
// is empty, it returns error.
func (q *partialResultQueue) peekLast() (*sppb.PartialResultSet, error) {
	if q.empty() {
		return nil, errEmptyQueue()
	}
	return q.q[(q.last+cap(q.q)-1)%cap(q.q)], nil
}

// push adds an item to the tail of partialResultQueue.
func (q *partialResultQueue) push(r *sppb.PartialResultSet) {
	if q.q == nil {
		q.q = make([]*sppb.PartialResultSet, 8 /* arbitrary */)
	}
	if q.n == cap(q.q) {
		buf := make([]*sppb.PartialResultSet, cap(q.q)*2)
		for i := 0; i < q.n; i++ {
			buf[i] = q.q[(q.first+i)%cap(q.q)]
		}
		q.q = buf
		q.first = 0
		q.last = q.n
	}
	q.q[q.last] = r
	q.last = (q.last + 1) % cap(q.q)
	q.n++
}

// pop removes an item from the head of partialResultQueue and returns it.
func (q *partialResultQueue) pop() *sppb.PartialResultSet {
	if q.n == 0 {
		return nil
	}
	r := q.q[q.first]
	q.q[q.first] = nil
	q.first = (q.first + 1) % cap(q.q)
	q.n--
	return r
}

// clear empties partialResultQueue.
func (q *partialResultQueue) clear() {
	*q = partialResultQueue{}
}

// dump retrieves all items from partialResultQueue and return them in a slice.
// It is used only in tests.
func (q *partialResultQueue) dump() []*sppb.PartialResultSet {
	var dq []*sppb.PartialResultSet
	for i := q.first; len(dq) < q.n; i = (i + 1) % cap(q.q) {
		dq = append(dq, q.q[i])
	}
	return dq
}

// resumableStreamDecoderState encodes resumableStreamDecoder's status. See also
// the comments for resumableStreamDecoder.Next.
type resumableStreamDecoderState int

const (
	unConnected         resumableStreamDecoderState = iota // 0
	queueingRetryable                                      // 1
	queueingUnretryable                                    // 2
	aborted                                                // 3
	finished                                               // 4
)

// resumableStreamDecoder provides a resumable interface for receiving
// sppb.PartialResultSet(s) from a given query wrapped by
// resumableStreamDecoder.rpc().
type resumableStreamDecoder struct {
	// state is the current status of resumableStreamDecoder, see also
	// the comments for resumableStreamDecoder.Next.
	state resumableStreamDecoderState

	// stateWitness when non-nil is called to observe state change,
	// used for testing.
	stateWitness func(resumableStreamDecoderState)

	// ctx is the caller's context, used for cancel/timeout Next().
	ctx context.Context

	// rpc is a factory of streamingReceiver, which might resume
	// a previous stream from the point encoded in restartToken.
	// rpc is always a wrapper of a Cloud Spanner query which is
	// resumable.
	rpc func(ctx context.Context, restartToken []byte) (streamingReceiver, error)

	// logger is the logger to use.
	logger *log.Logger

	// stream is the current RPC streaming receiver.
	stream streamingReceiver

	// q buffers received yet undecoded partial results.
	q partialResultQueue

	// bytesBetweenResumeTokens is the proxy of the byte size of
	// PartialResultSets being queued between two resume tokens. Once
	// bytesBetweenResumeTokens is greater than maxBytesBetweenResumeTokens,
	// resumableStreamDecoder goes into queueingUnretryable state.
	bytesBetweenResumeTokens int32

	// maxBytesBetweenResumeTokens is the max number of bytes that can be
	// buffered between two resume tokens. It is always copied from the global
	// maxBytesBetweenResumeTokens atomically.
	maxBytesBetweenResumeTokens int32

	// np is the next sppb.PartialResultSet ready to be returned
	// to caller of resumableStreamDecoder.Get().
	np *sppb.PartialResultSet

	// resumeToken stores the resume token that resumableStreamDecoder has
	// last revealed to caller.
	resumeToken []byte

	// err is the last error resumableStreamDecoder has encountered so far.
	err error

	// backoff is used for the retry settings
	backoff gax.Backoff
}

// newResumableStreamDecoder creates a new resumeableStreamDecoder instance.
// Parameter rpc should be a function that creates a new stream beginning at the
// restartToken if non-nil.
func newResumableStreamDecoder(ctx context.Context, logger *log.Logger, rpc func(ct context.Context, restartToken []byte) (streamingReceiver, error)) *resumableStreamDecoder {
	return &resumableStreamDecoder{
		ctx:                         ctx,
		logger:                      logger,
		rpc:                         rpc,
		maxBytesBetweenResumeTokens: atomic.LoadInt32(&maxBytesBetweenResumeTokens),
		backoff:                     DefaultRetryBackoff,
	}
}

// changeState fulfills state transition for resumableStateDecoder.
func (d *resumableStreamDecoder) changeState(target resumableStreamDecoderState) {
	if d.state == queueingRetryable && d.state != target {
		// Reset bytesBetweenResumeTokens because it is only meaningful/changed
		// under queueingRetryable state.
		d.bytesBetweenResumeTokens = 0
	}
	d.state = target
	if d.stateWitness != nil {
		d.stateWitness(target)
	}
}

// isNewResumeToken returns if the observed resume token is different from
// the one returned from server last time.
func (d *resumableStreamDecoder) isNewResumeToken(rt []byte) bool {
	if rt == nil {
		return false
	}
	if bytes.Equal(rt, d.resumeToken) {
		return false
	}
	return true
}

// Next advances to the next available partial result set.  If error or no
// more, returns false, call Err to determine if an error was encountered.
// The following diagram illustrates the state machine of resumableStreamDecoder
// that Next() implements. Note that state transition can be only triggered by
// RPC activities.
/*
        rpc() fails retryable
      +---------+
      |         |    rpc() fails unretryable/ctx timeouts or cancelled
      |         |   +------------------------------------------------+
      |         |   |                                                |
      |         v   |                                                v
      |     +---+---+---+                       +--------+    +------+--+
      +-----+unConnected|                       |finished|    | aborted |<----+
            |           |                       ++-----+-+    +------+--+     |
            +---+----+--+                        ^     ^             ^        |
                |    ^                           |     |             |        |
                |    |                           |     |     recv() fails     |
                |    |                           |     |             |        |
                |    |recv() fails retryable     |     |             |        |
                |    |with valid ctx             |     |             |        |
                |    |                           |     |             |        |
      rpc() succeeds |   +-----------------------+     |             |        |
                |    |   |         recv EOF         recv EOF         |        |
                |    |   |                             |             |        |
                v    |   |     Queue size exceeds      |             |        |
            +---+----+---+----+threshold       +-------+-----------+ |        |
+---------->+                 +--------------->+                   +-+        |
|           |queueingRetryable|                |queueingUnretryable|          |
|           |                 +<---------------+                   |          |
|           +---+----------+--+ pop() returns  +--+----+-----------+          |
|               |          |    resume token      |    ^                      |
|               |          |                      |    |                      |
|               |          |                      |    |                      |
+---------------+          |                      |    |                      |
   recv() succeeds         |                      +----+                      |
                           |                      recv() succeeds             |
                           |                                                  |
                           |                                                  |
                           |                                                  |
                           |                                                  |
                           |                                                  |
                           +--------------------------------------------------+
                                               recv() fails unretryable

*/
var (
	// maxBytesBetweenResumeTokens is the maximum amount of bytes that
	// resumableStreamDecoder in queueingRetryable state can use to queue
	// PartialResultSets before getting into queueingUnretryable state.
	maxBytesBetweenResumeTokens = int32(128 * 1024 * 1024)
)

func (d *resumableStreamDecoder) next() bool {
	retryer := gax.OnCodes([]codes.Code{codes.Unavailable, codes.Internal}, d.backoff)
	for {
		switch d.state {
		case unConnected:
			// If no gRPC stream is available, try to initiate one.
			d.stream, d.err = d.rpc(d.ctx, d.resumeToken)
			if d.err == nil {
				d.changeState(queueingRetryable)
				continue
			}
			delay, shouldRetry := retryer.Retry(d.err)
			if !shouldRetry {
				d.changeState(aborted)
				continue
			}
			trace.TracePrintf(d.ctx, nil, "Backing off stream read for %s", delay)
			if err := gax.Sleep(d.ctx, delay); err == nil {
				// Be explicit about state transition, although the
				// state doesn't actually change. State transition
				// will be triggered only by RPC activity, regardless of
				// whether there is an actual state change or not.
				d.changeState(unConnected)
			} else {
				d.err = err
				d.changeState(aborted)
			}
			continue

		case queueingRetryable:
			fallthrough
		case queueingUnretryable:
			// Receiving queue is not empty.
			last, err := d.q.peekLast()
			if err != nil {
				// Only the case that receiving queue is empty could cause
				// peekLast to return error and in such case, we should try to
				// receive from stream.
				d.tryRecv(retryer)
				continue
			}
			if d.isNewResumeToken(last.ResumeToken) {
				// Got new resume token, return buffered sppb.PartialResultSets
				// to caller.
				d.np = d.q.pop()
				if d.q.empty() {
					d.bytesBetweenResumeTokens = 0
					// The new resume token was just popped out from queue,
					// record it.
					d.resumeToken = d.np.ResumeToken
					d.changeState(queueingRetryable)
				}
				return true
			}
			if d.bytesBetweenResumeTokens >= d.maxBytesBetweenResumeTokens && d.state == queueingRetryable {
				d.changeState(queueingUnretryable)
				continue
			}
			if d.state == queueingUnretryable {
				// When there is no resume token observed, only yield
				// sppb.PartialResultSets to caller under queueingUnretryable
				// state.
				d.np = d.q.pop()
				return true
			}
			// Needs to receive more from gRPC stream till a new resume token
			// is observed.
			d.tryRecv(retryer)
			continue
		case aborted:
			// Discard all pending items because none of them should be yield
			// to caller.
			d.q.clear()
			return false
		case finished:
			// If query has finished, check if there are still buffered messages.
			if d.q.empty() {
				// No buffered PartialResultSet.
				return false
			}
			// Although query has finished, there are still buffered
			// PartialResultSets.
			d.np = d.q.pop()
			return true

		default:
			logf(d.logger, "Unexpected resumableStreamDecoder.state: %v", d.state)
			return false
		}
	}
}

// tryRecv attempts to receive a PartialResultSet from gRPC stream.
func (d *resumableStreamDecoder) tryRecv(retryer gax.Retryer) {
	var res *sppb.PartialResultSet
	res, d.err = d.stream.Recv()
	if d.err == nil {
		d.q.push(res)
		if d.state == queueingRetryable && !d.isNewResumeToken(res.ResumeToken) {
			d.bytesBetweenResumeTokens += int32(proto.Size(res))
		}
		d.changeState(d.state)
		return
	}
	if d.err == io.EOF {
		d.err = nil
		d.changeState(finished)
		return
	}
	delay, shouldRetry := retryer.Retry(d.err)
	if !shouldRetry || d.state != queueingRetryable {
		d.changeState(aborted)
		return
	}
	if err := gax.Sleep(d.ctx, delay); err != nil {
		d.err = err
		d.changeState(aborted)
		return
	}
	// Clear error and retry the stream.
	d.err = nil
	// Discard all queue items (none have resume tokens).
	d.q.clear()
	d.stream = nil
	d.changeState(unConnected)
}

// get returns the most recent PartialResultSet generated by a call to next.
func (d *resumableStreamDecoder) get() *sppb.PartialResultSet {
	return d.np
}

// lastErr returns the last non-EOF error encountered.
func (d *resumableStreamDecoder) lastErr() error {
	return d.err
}

// partialResultSetDecoder assembles PartialResultSet(s) into Cloud Spanner
// Rows.
type partialResultSetDecoder struct {
	row     Row
	tx      *sppb.Transaction
	chunked bool // if true, next value should be merged with last values
	// entry.
	ts time.Time // read timestamp
}

// yield checks we have a complete row, and if so returns it.  A row is not
// complete if it doesn't have enough columns, or if this is a chunked response
// and there are no further values to process.
func (p *partialResultSetDecoder) yield(chunked, last bool) *Row {
	if len(p.row.vals) == len(p.row.fields) && (!chunked || !last) {
		// When partialResultSetDecoder gets enough number of Column values.
		// There are two cases that a new Row should be yield:
		//
		//   1. The incoming PartialResultSet is not chunked;
		//   2. The incoming PartialResultSet is chunked, but the
		//      proto3.Value being merged is not the last one in
		//      the PartialResultSet.
		//
		// Use a fresh Row to simplify clients that want to use yielded results
		// after the next row is retrieved. Note that fields is never changed
		// so it doesn't need to be copied.
		fresh := Row{
			fields: p.row.fields,
			vals:   make([]*proto3.Value, len(p.row.vals)),
		}
		copy(fresh.vals, p.row.vals)
		p.row.vals = p.row.vals[:0] // empty and reuse slice
		return &fresh
	}
	return nil
}

// yieldTx returns transaction information via caller supplied callback.
func errChunkedEmptyRow() error {
	return spannerErrorf(codes.FailedPrecondition, "got invalid chunked PartialResultSet with empty Row")
}

// add tries to merge a new PartialResultSet into buffered Row. It returns any
// rows that have been completed as a result.
func (p *partialResultSetDecoder) add(r *sppb.PartialResultSet) ([]*Row, error) {
	var rows []*Row
	if r.Metadata != nil {
		// Metadata should only be returned in the first result.
		if p.row.fields == nil {
			p.row.fields = r.Metadata.RowType.Fields
		}
		if p.tx == nil && r.Metadata.Transaction != nil {
			p.tx = r.Metadata.Transaction
			if p.tx.ReadTimestamp != nil {
				p.ts = time.Unix(p.tx.ReadTimestamp.Seconds, int64(p.tx.ReadTimestamp.Nanos))
			}
		}
	}
	if len(r.Values) == 0 {
		return nil, nil
	}
	if p.chunked {
		p.chunked = false
		// Try to merge first value in r.Values into uncompleted row.
		last := len(p.row.vals) - 1
		if last < 0 { // sanity check
			return nil, errChunkedEmptyRow()
		}
		var err error
		// If p is chunked, then we should always try to merge p.last with
		// r.first.
		if p.row.vals[last], err = p.merge(p.row.vals[last], r.Values[0]); err != nil {
			return nil, err
		}
		r.Values = r.Values[1:]
		// Merge is done, try to yield a complete Row.
		if row := p.yield(r.ChunkedValue, len(r.Values) == 0); row != nil {
			rows = append(rows, row)
		}
	}
	for i, v := range r.Values {
		// The rest values in r can be appened into p directly.
		p.row.vals = append(p.row.vals, v)
		// Again, check to see if a complete Row can be yielded because of the
		// newly added value.
		if row := p.yield(r.ChunkedValue, i == len(r.Values)-1); row != nil {
			rows = append(rows, row)
		}
	}
	if r.ChunkedValue {
		// After dealing with all values in r, if r is chunked then p must be
		// also chunked.
		p.chunked = true
	}
	return rows, nil
}

// isMergeable returns if a protobuf Value can be potentially merged with other
// protobuf Values.
func (p *partialResultSetDecoder) isMergeable(a *proto3.Value) bool {
	switch a.Kind.(type) {
	case *proto3.Value_StringValue:
		return true
	case *proto3.Value_ListValue:
		return true
	default:
		return false
	}
}

// errIncompatibleMergeTypes returns error for incompatible protobuf types that
// cannot be merged by partialResultSetDecoder.
func errIncompatibleMergeTypes(a, b *proto3.Value) error {
	return spannerErrorf(codes.FailedPrecondition, "incompatible type in chunked PartialResultSet. expected (%T), got (%T)", a.Kind, b.Kind)
}

// errUnsupportedMergeType returns error for protobuf type that cannot be merged
// to other protobufs.
func errUnsupportedMergeType(a *proto3.Value) error {
	return spannerErrorf(codes.FailedPrecondition, "unsupported type merge (%T)", a.Kind)
}

// merge tries to combine two protobuf Values if possible.
func (p *partialResultSetDecoder) merge(a, b *proto3.Value) (*proto3.Value, error) {
	var err error
	typeErr := errIncompatibleMergeTypes(a, b)
	switch t := a.Kind.(type) {
	case *proto3.Value_StringValue:
		s, ok := b.Kind.(*proto3.Value_StringValue)
		if !ok {
			return nil, typeErr
		}
		return &proto3.Value{
			Kind: &proto3.Value_StringValue{StringValue: t.StringValue + s.StringValue},
		}, nil
	case *proto3.Value_ListValue:
		l, ok := b.Kind.(*proto3.Value_ListValue)
		if !ok {
			return nil, typeErr
		}
		if l.ListValue == nil || len(l.ListValue.Values) <= 0 {
			// b is an empty list, just return a.
			return a, nil
		}
		if t.ListValue == nil || len(t.ListValue.Values) <= 0 {
			// a is an empty list, just return b.
			return b, nil
		}
		if la := len(t.ListValue.Values) - 1; p.isMergeable(t.ListValue.Values[la]) {
			// When the last item in a is of type String, List or Struct
			// (encoded into List by Cloud Spanner), try to Merge last item in
			// a and first item in b.
			t.ListValue.Values[la], err = p.merge(t.ListValue.Values[la], l.ListValue.Values[0])
			if err != nil {
				return nil, err
			}
			l.ListValue.Values = l.ListValue.Values[1:]
		}
		return &proto3.Value{
			Kind: &proto3.Value_ListValue{
				ListValue: &proto3.ListValue{
					Values: append(t.ListValue.Values, l.ListValue.Values...),
				},
			},
		}, nil
	default:
		return nil, errUnsupportedMergeType(a)
	}

}

// Done returns if partialResultSetDecoder has already done with all buffered
// values.
func (p *partialResultSetDecoder) done() bool {
	// There is no explicit end of stream marker, but ending part way through a
	// row is obviously bad, or ending with the last column still awaiting
	// completion.
	return len(p.row.vals) == 0 && !p.chunked
}