// getQuerySplitToKeyRangePartFunc returns a function to use with scatterConn.SplitQueryV2
// that converts the given QuerySplit to a SplitQueryResponse_Part message whose KeyRangePart field
// is set.
func getQuerySplitToKeyRangePartFunc(
	keyspace string,
	shardReferenceByName map[string]*topodatapb.ShardReference) func(
	querySplit *querytypes.QuerySplit, shard string) (*vtgatepb.SplitQueryResponse_Part, error) {

	return func(
		querySplit *querytypes.QuerySplit, shard string) (*vtgatepb.SplitQueryResponse_Part, error) {
		// TODO(erez): Assert that shardReferenceByName contains an entry for 'shard'.
		// Keyrange can be nil for the shard (e.g. for single-sharded keyspaces during resharding).
		// In this case we append an empty keyrange that represents the entire keyspace.
		keyranges := []*topodatapb.KeyRange{{Start: []byte{}, End: []byte{}}}
		if shardReferenceByName[shard].KeyRange != nil {
			keyranges = []*topodatapb.KeyRange{shardReferenceByName[shard].KeyRange}
		}
		bindVars, err := querytypes.BindVariablesToProto3(querySplit.BindVariables)
		if err != nil {
			return nil, err
		}
		return &vtgatepb.SplitQueryResponse_Part{
			Query: &querypb.BoundQuery{
				Sql:           querySplit.Sql,
				BindVariables: bindVars,
			},
			KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{
				Keyspace:  keyspace,
				KeyRanges: keyranges,
			},
			Size: querySplit.RowCount,
		}, nil
	}
}

// StreamExecute is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) StreamExecute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (<-chan *sqltypes.Result, tabletconn.ErrFunc, error) {
	bv, err := querytypes.BindVariablesToProto3(bindVars)
	if err != nil {
		return nil, nil, err
	}
	bindVars, err = querytypes.Proto3ToBindVariables(bv)
	if err != nil {
		return nil, nil, err
	}
	result := make(chan *sqltypes.Result, 10)
	var finalErr error

	go func() {
		finalErr = itc.tablet.qsc.QueryService().StreamExecute(ctx, &querypb.Target{
			Keyspace:   itc.tablet.keyspace,
			Shard:      itc.tablet.shard,
			TabletType: itc.tablet.tabletType,
		}, query, bindVars, 0, func(reply *sqltypes.Result) error {
			// We need to deep-copy the reply before returning,
			// because the underlying buffers are reused.
			result <- reply.Copy()
			return nil
		})

		// the client will only access finalErr after the
		// channel is closed, and then it's already set.
		close(result)
	}()

	return result, func() error {
		return tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(finalErr))
	}, nil
}

// StreamExecute is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) StreamExecute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]interface{}, options *querypb.ExecuteOptions) (sqltypes.ResultStream, error) {
	bv, err := querytypes.BindVariablesToProto3(bindVars)
	if err != nil {
		return nil, err
	}
	bindVars, err = querytypes.Proto3ToBindVariables(bv)
	if err != nil {
		return nil, err
	}
	result := make(chan *sqltypes.Result, 10)
	var finalErr error

	go func() {
		finalErr = itc.tablet.qsc.QueryService().StreamExecute(ctx, target, query, bindVars, options, func(reply *sqltypes.Result) error {
			// We need to deep-copy the reply before returning,
			// because the underlying buffers are reused.
			result <- reply.Copy()
			return nil
		})
		finalErr = tabletconn.TabletErrorFromGRPC(vterrors.ToGRPCError(finalErr))

		// the client will only access finalErr after the
		// channel is closed, and then it's already set.
		close(result)
	}()

	return &streamExecuteAdapter{result, &finalErr}, nil
}

// SplitQueryKeyRange scatters a SplitQuery request to all shards. For a set of
// splits received from a shard, it construct a KeyRange queries by
// appending that shard's keyrange to the splits. Aggregates all splits across
// all shards in no specific order and returns.
func (stc *ScatterConn) SplitQueryKeyRange(ctx context.Context, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64, keyRangeByShard map[string]*topodatapb.KeyRange, keyspace string) ([]*vtgatepb.SplitQueryResponse_Part, error) {
	tabletType := topodatapb.TabletType_RDONLY
	actionFunc := func(shard string, transactionID int64, results chan<- interface{}) error {
		// Get all splits from this shard
		queries, err := stc.gateway.SplitQuery(ctx, keyspace, shard, tabletType, sql, bindVariables, splitColumn, splitCount)
		if err != nil {
			return err
		}
		// Append the keyrange for this shard to all the splits received,
		// if keyrange is nil for the shard (e.g. for single-sharded keyspaces during resharding),
		// append empty keyrange to represent the entire keyspace.
		keyranges := []*topodatapb.KeyRange{{Start: []byte{}, End: []byte{}}}
		if keyRangeByShard[shard] != nil {
			keyranges = []*topodatapb.KeyRange{keyRangeByShard[shard]}
		}
		splits := make([]*vtgatepb.SplitQueryResponse_Part, len(queries))
		for i, query := range queries {
			q, err := querytypes.BindVariablesToProto3(query.BindVariables)
			if err != nil {
				return err
			}
			splits[i] = &vtgatepb.SplitQueryResponse_Part{
				Query: &querypb.BoundQuery{
					Sql:           query.Sql,
					BindVariables: q,
				},
				KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{
					Keyspace:  keyspace,
					KeyRanges: keyranges,
				},
				Size: query.RowCount,
			}
		}
		// Push all the splits from this shard to results channel
		results <- splits
		return nil
	}

	shards := []string{}
	for shard := range keyRangeByShard {
		shards = append(shards, shard)
	}
	allSplits, allErrors := stc.multiGo(ctx, "SplitQuery", keyspace, shards, tabletType, NewSafeSession(&vtgatepb.Session{}), false, actionFunc)
	splits := []*vtgatepb.SplitQueryResponse_Part{}
	for s := range allSplits {
		splits = append(splits, s.([]*vtgatepb.SplitQueryResponse_Part)...)
	}
	if allErrors.HasErrors() {
		err := allErrors.AggrError(stc.aggregateErrors)
		return nil, err
	}
	return splits, nil
}

// SplitQueryCustomSharding scatters a SplitQuery request to all
// shards. For a set of splits received from a shard, it construct a
// KeyRange queries by appending that shard's name to the
// splits. Aggregates all splits across all shards in no specific
// order and returns.
func (stc *ScatterConn) SplitQueryCustomSharding(ctx context.Context, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64, shards []string, keyspace string) ([]*vtgatepb.SplitQueryResponse_Part, error) {
	tabletType := topodatapb.TabletType_RDONLY

	// mu protects allSplits
	var mu sync.Mutex
	var allSplits []*vtgatepb.SplitQueryResponse_Part

	actionFunc := func(target *querypb.Target) error {
		// Get all splits from this shard
		query := querytypes.BoundQuery{
			Sql:           sql,
			BindVariables: bindVariables,
		}
		queries, err := stc.gateway.SplitQuery(ctx, target, query, splitColumn, splitCount)
		if err != nil {
			return err
		}
		// Use the shards list for all the splits received
		shards := []string{target.Shard}
		splits := make([]*vtgatepb.SplitQueryResponse_Part, len(queries))
		for i, query := range queries {
			q, err := querytypes.BindVariablesToProto3(query.BindVariables)
			if err != nil {
				return err
			}
			splits[i] = &vtgatepb.SplitQueryResponse_Part{
				Query: &querypb.BoundQuery{
					Sql:           query.Sql,
					BindVariables: q,
				},
				ShardPart: &vtgatepb.SplitQueryResponse_ShardPart{
					Keyspace: keyspace,
					Shards:   shards,
				},
				Size: query.RowCount,
			}
		}

		// aggregate splits
		mu.Lock()
		defer mu.Unlock()
		allSplits = append(allSplits, splits...)
		return nil
	}
	allErrors := stc.multiGo(ctx, "SplitQuery", keyspace, shards, tabletType, actionFunc)
	if allErrors.HasErrors() {
		return nil, allErrors.AggrError(stc.aggregateErrors)
	}
	// See the comment for the analogues line in SplitQueryKeyRange for
	// the motivation for shuffling.
	shuffleQueryParts(allSplits)
	return allSplits, nil
}

// Execute is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) Execute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]interface{}, transactionID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
	bv, err := querytypes.BindVariablesToProto3(bindVars)
	if err != nil {
		return nil, err
	}
	bindVars, err = querytypes.Proto3ToBindVariables(bv)
	if err != nil {
		return nil, err
	}
	reply, err := itc.tablet.qsc.QueryService().Execute(ctx, target, query, bindVars, transactionID, options)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(vterrors.ToGRPCError(err))
	}
	return reply, nil
}

// SplitQueryCustomSharding scatters a SplitQuery request to all
// shards. For a set of splits received from a shard, it construct a
// KeyRange queries by appending that shard's name to the
// splits. Aggregates all splits across all shards in no specific
// order and returns.
func (stc *ScatterConn) SplitQueryCustomSharding(ctx context.Context, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64, shards []string, keyspace string) ([]*vtgatepb.SplitQueryResponse_Part, error) {
	tabletType := topodatapb.TabletType_RDONLY
	actionFunc := func(shard string, transactionID int64, results chan<- interface{}) error {
		// Get all splits from this shard
		queries, err := stc.gateway.SplitQuery(ctx, keyspace, shard, tabletType, sql, bindVariables, splitColumn, splitCount)
		if err != nil {
			return err
		}
		// Use the shards list for all the splits received
		shards := []string{shard}
		splits := make([]*vtgatepb.SplitQueryResponse_Part, len(queries))
		for i, query := range queries {
			q, err := querytypes.BindVariablesToProto3(query.BindVariables)
			if err != nil {
				return err
			}
			splits[i] = &vtgatepb.SplitQueryResponse_Part{
				Query: &querypb.BoundQuery{
					Sql:           query.Sql,
					BindVariables: q,
				},
				ShardPart: &vtgatepb.SplitQueryResponse_ShardPart{
					Keyspace: keyspace,
					Shards:   shards,
				},
				Size: query.RowCount,
			}
		}
		// Push all the splits from this shard to results channel
		results <- splits
		return nil
	}

	allSplits, allErrors := stc.multiGo(ctx, "SplitQuery", keyspace, shards, tabletType, NewSafeSession(&vtgatepb.Session{}), false, actionFunc)
	splits := []*vtgatepb.SplitQueryResponse_Part{}
	for s := range allSplits {
		splits = append(splits, s.([]*vtgatepb.SplitQueryResponse_Part)...)
	}
	if allErrors.HasErrors() {
		err := allErrors.AggrError(stc.aggregateErrors)
		return nil, err
	}
	return splits, nil
}

// Execute is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) Execute(ctx context.Context, query string, bindVars map[string]interface{}, transactionID int64) (*sqltypes.Result, error) {
	bv, err := querytypes.BindVariablesToProto3(bindVars)
	if err != nil {
		return nil, err
	}
	bindVars, err = querytypes.Proto3ToBindVariables(bv)
	if err != nil {
		return nil, err
	}
	reply, err := itc.tablet.qsc.QueryService().Execute(ctx, &querypb.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, query, bindVars, 0, transactionID)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return reply, nil
}

func (c *echoClient) SplitQuery(ctx context.Context, keyspace string, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64) ([]*vtgatepb.SplitQueryResponse_Part, error) {
	if strings.HasPrefix(sql, EchoPrefix) {
		bv, err := querytypes.BindVariablesToProto3(bindVariables)
		if err != nil {
			return nil, err
		}
		return []*vtgatepb.SplitQueryResponse_Part{
			{
				Query: &querypb.BoundQuery{
					Sql:           fmt.Sprintf("%v:%v:%v", sql, splitColumn, splitCount),
					BindVariables: bv,
				},
				KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{
					Keyspace: keyspace,
				},
			},
		}, nil
	}
	return c.fallback.SplitQuery(ctx, sql, keyspace, bindVariables, splitColumn, splitCount)
}

// ExecuteBatch is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) ExecuteBatch(ctx context.Context, target *querypb.Target, queries []querytypes.BoundQuery, asTransaction bool, transactionID int64, options *querypb.ExecuteOptions) ([]sqltypes.Result, error) {
	q := make([]querytypes.BoundQuery, len(queries))
	for i, query := range queries {
		bv, err := querytypes.BindVariablesToProto3(query.BindVariables)
		if err != nil {
			return nil, err
		}
		bindVars, err := querytypes.Proto3ToBindVariables(bv)
		if err != nil {
			return nil, err
		}
		q[i].Sql = query.Sql
		q[i].BindVariables = bindVars
	}
	results, err := itc.tablet.qsc.QueryService().ExecuteBatch(ctx, target, q, asTransaction, transactionID, options)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(vterrors.ToGRPCError(err))
	}
	return results, nil
}

// getQuerySplitToShardPartFunc returns a function to use with scatterConn.SplitQueryV2
// that converts the given QuerySplit to a SplitQueryResponse_Part message whose ShardPart field
// is set.
func getQuerySplitToShardPartFunc(keyspace string) func(
	querySplit *querytypes.QuerySplit, shard string) (*vtgatepb.SplitQueryResponse_Part, error) {

	return func(
		querySplit *querytypes.QuerySplit, shard string) (*vtgatepb.SplitQueryResponse_Part, error) {
		bindVars, err := querytypes.BindVariablesToProto3(querySplit.BindVariables)
		if err != nil {
			return nil, err
		}
		return &vtgatepb.SplitQueryResponse_Part{
			Query: &querypb.BoundQuery{
				Sql:           querySplit.Sql,
				BindVariables: bindVars,
			},
			ShardPart: &vtgatepb.SplitQueryResponse_ShardPart{
				Keyspace: keyspace,
				Shards:   []string{shard},
			},
			Size: querySplit.RowCount,
		}, nil
	}
}

// ExecuteBatch is part of tabletconn.TabletConn
// We need to copy the bind variables as tablet server will change them.
func (itc *internalTabletConn) ExecuteBatch(ctx context.Context, queries []querytypes.BoundQuery, asTransaction bool, transactionID int64) ([]sqltypes.Result, error) {
	q := make([]querytypes.BoundQuery, len(queries))
	for i, query := range queries {
		bv, err := querytypes.BindVariablesToProto3(query.BindVariables)
		if err != nil {
			return nil, err
		}
		bindVars, err := querytypes.Proto3ToBindVariables(bv)
		if err != nil {
			return nil, err
		}
		q[i].Sql = query.Sql
		q[i].BindVariables = bindVars
	}
	results, err := itc.tablet.qsc.QueryService().ExecuteBatch(ctx, &querypb.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, q, 0, asTransaction, transactionID)
	if err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return results, nil
}

// SplitQueryKeyRange scatters a SplitQuery request to all shards. For a set of
// splits received from a shard, it construct a KeyRange queries by
// appending that shard's keyrange to the splits. Aggregates all splits across
// all shards in no specific order and returns.
func (stc *ScatterConn) SplitQueryKeyRange(ctx context.Context, sql string, bindVariables map[string]interface{}, splitColumn string, splitCount int64, keyRangeByShard map[string]*topodatapb.KeyRange, keyspace string) ([]*vtgatepb.SplitQueryResponse_Part, error) {
	tabletType := topodatapb.TabletType_RDONLY

	// mu protects allSplits
	var mu sync.Mutex
	var allSplits []*vtgatepb.SplitQueryResponse_Part

	actionFunc := func(shard string, transactionID int64) error {
		// Get all splits from this shard
		queries, err := stc.gateway.SplitQuery(ctx, keyspace, shard, tabletType, sql, bindVariables, splitColumn, splitCount)
		if err != nil {
			return err
		}
		// Append the keyrange for this shard to all the splits received,
		// if keyrange is nil for the shard (e.g. for single-sharded keyspaces during resharding),
		// append empty keyrange to represent the entire keyspace.
		keyranges := []*topodatapb.KeyRange{{Start: []byte{}, End: []byte{}}}
		if keyRangeByShard[shard] != nil {
			keyranges = []*topodatapb.KeyRange{keyRangeByShard[shard]}
		}
		splits := make([]*vtgatepb.SplitQueryResponse_Part, len(queries))
		for i, query := range queries {
			q, err := querytypes.BindVariablesToProto3(query.BindVariables)
			if err != nil {
				return err
			}
			splits[i] = &vtgatepb.SplitQueryResponse_Part{
				Query: &querypb.BoundQuery{
					Sql:           query.Sql,
					BindVariables: q,
				},
				KeyRangePart: &vtgatepb.SplitQueryResponse_KeyRangePart{
					Keyspace:  keyspace,
					KeyRanges: keyranges,
				},
				Size: query.RowCount,
			}
		}

		// aggregate splits
		mu.Lock()
		defer mu.Unlock()
		allSplits = append(allSplits, splits...)
		return nil
	}

	shards := []string{}
	for shard := range keyRangeByShard {
		shards = append(shards, shard)
	}
	allErrors := stc.multiGo(ctx, "SplitQuery", keyspace, shards, tabletType, NewSafeSession(&vtgatepb.Session{}), false, actionFunc)
	if allErrors.HasErrors() {
		return nil, allErrors.AggrError(stc.aggregateErrors)
	}
	// We shuffle the query-parts here. External frameworks like MapReduce may
	// "deal" these jobs to workers in the order they are in the list. Without
	// shuffling workers can be very unevenly distributed among
	// the shards they query. E.g. all workers will first query the first shard,
	// then most of them to the second shard, etc, which results with uneven
	// load balancing among shards.
	shuffleQueryParts(allSplits)
	return allSplits, nil
}