// 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 []tproto.BoundQuery, asTransaction bool, transactionID int64) (*tproto.QueryResultList, error) {
	q := make([]tproto.BoundQuery, len(queries))
	for i, query := range queries {
		bv, err := tproto.BindVariablesToProto3(query.BindVariables)
		if err != nil {
			return nil, err
		}
		bindVars, err := tproto.Proto3ToBindVariables(bv)
		if err != nil {
			return nil, err
		}
		q[i].Sql = query.Sql
		q[i].BindVariables = bindVars
	}
	reply := &tproto.QueryResultList{}
	if err := itc.tablet.qsc.QueryService().ExecuteBatch(ctx, &querypb.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.QueryList{
		Queries:       q,
		AsTransaction: asTransaction,
		TransactionId: transactionID,
	}, reply); err != nil {
		return nil, tabletconn.TabletErrorFromGRPC(tabletserver.ToGRPCError(err))
	}
	return reply, 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 *mproto.QueryResult, tabletconn.ErrFunc, error) {
	bindVars = tproto.Proto3ToBindVariables(tproto.BindVariablesToProto3(bindVars))
	result := make(chan *mproto.QueryResult, 10)
	var finalErr error

	go func() {
		finalErr = itc.tablet.qsc.QueryService().StreamExecute(ctx, &pbq.Target{
			Keyspace:   itc.tablet.keyspace,
			Shard:      itc.tablet.shard,
			TabletType: itc.tablet.tabletType,
		}, &tproto.Query{
			Sql:           query,
			BindVariables: bindVars,
			TransactionId: transactionID,
		}, func(reply *mproto.QueryResult) error {
			result <- reply
			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
}

// 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 int, keyRangeByShard map[string]*pb.KeyRange, keyspace string) ([]*pbg.SplitQueryResponse_Part, error) {
	tabletType := pb.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
		keyranges := []*pb.KeyRange{keyRangeByShard[shard]}
		splits := make([]*pbg.SplitQueryResponse_Part, len(queries))
		for i, query := range queries {
			q, err := tproto.BindVariablesToProto3(query.Query.BindVariables)
			if err != nil {
				return err
			}
			splits[i] = &pbg.SplitQueryResponse_Part{
				Query: &pbq.BoundQuery{
					Sql:           query.Query.Sql,
					BindVariables: q,
				},
				KeyRangePart: &pbg.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(&pbg.Session{}), false, actionFunc)
	splits := []*pbg.SplitQueryResponse_Part{}
	for s := range allSplits {
		splits = append(splits, s.([]*pbg.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) (*mproto.QueryResult, error) {
	bindVars = tproto.Proto3ToBindVariables(tproto.BindVariablesToProto3(bindVars))
	reply := &mproto.QueryResult{}
	if err := itc.tablet.qsc.QueryService().Execute(ctx, &pbq.Target{
		Keyspace:   itc.tablet.keyspace,
		Shard:      itc.tablet.shard,
		TabletType: itc.tablet.tabletType,
	}, &tproto.Query{
		Sql:           query,
		BindVariables: bindVars,
		TransactionId: transactionID,
	}, reply); 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 int) ([]*pbg.SplitQueryResponse_Part, error) {
	if strings.HasPrefix(sql, EchoPrefix) {
		return []*pbg.SplitQueryResponse_Part{
			&pbg.SplitQueryResponse_Part{
				Query: &pbq.BoundQuery{
					Sql:           fmt.Sprintf("%v:%v:%v", sql, splitColumn, splitCount),
					BindVariables: tproto.BindVariablesToProto3(bindVariables),
				},
				KeyRangePart: &pbg.SplitQueryResponse_KeyRangePart{
					Keyspace: keyspace,
				},
			},
		}, nil
	}
	return c.fallback.SplitQuery(ctx, sql, keyspace, bindVariables, splitColumn, splitCount)
}

// 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 int, 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 := tproto.BindVariablesToProto3(query.Query.BindVariables)
			if err != nil {
				return err
			}
			splits[i] = &vtgatepb.SplitQueryResponse_Part{
				Query: &querypb.BoundQuery{
					Sql:           query.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
}

// 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 := tproto.BindVariablesToProto3(bindVars)
	if err != nil {
		return nil, nil, err
	}
	bindVars, err = tproto.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,
		}, &tproto.Query{
			Sql:           query,
			BindVariables: bindVars,
			TransactionId: transactionID,
		}, 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
}