// getConn reuses an existing connection if possible. Otherwise
// it returns a connection which it will save for future reuse.
// If it returns an error, retry will tell you if getConn can be retried.
// If the context has a deadline and exceeded, it returns error and no-retry immediately.
func (sdc *ShardConn) getConn(ctx context.Context) (conn tabletconn.TabletConn, endPoint topo.EndPoint, err error, retry bool) {
	sdc.mu.Lock()
	defer sdc.mu.Unlock()

	// fail-fast if deadline exceeded
	deadline, ok := ctx.Deadline()
	if ok {
		if time.Now().After(deadline) {
			return nil, topo.EndPoint{}, tabletconn.OperationalError("vttablet: deadline exceeded"), false
		}
	}

	if sdc.conn != nil {
		return sdc.conn, sdc.conn.EndPoint(), nil, false
	}

	endPoint, err = sdc.balancer.Get()
	if err != nil {
		return nil, topo.EndPoint{}, err, false
	}
	conn, err = tabletconn.GetDialer()(ctx, endPoint, sdc.keyspace, sdc.shard, sdc.timeout)
	if err != nil {
		sdc.balancer.MarkDown(endPoint.Uid, err.Error())
		return nil, endPoint, err, true
	}
	sdc.conn = conn
	return sdc.conn, endPoint, nil, false
}

func tabletError(err error) error {
	if err == nil {
		return nil
	}
	// TODO(aaijazi): tabletconn is in an intermediate state right now, where application errors
	// can be returned as rpcplus.ServerError or vterrors.VitessError. Soon, it will be standardized
	// to only VitessError.
	if ve, ok := err.(*vterrors.VitessError); ok {
		return tabletErrorFromVitessError(ve)
	}
	if _, ok := err.(rpcplus.ServerError); ok {
		var code int
		errStr := err.Error()
		switch {
		case strings.Contains(errStr, "fatal: "):
			code = tabletconn.ERR_FATAL
		case strings.Contains(errStr, "retry: "):
			code = tabletconn.ERR_RETRY
		case strings.Contains(errStr, "tx_pool_full: "):
			code = tabletconn.ERR_TX_POOL_FULL
		case strings.Contains(errStr, "not_in_tx: "):
			code = tabletconn.ERR_NOT_IN_TX
		default:
			code = tabletconn.ERR_NORMAL
		}
		return &tabletconn.ServerError{Code: code, Err: fmt.Sprintf("vttablet: %v", err)}
	}
	if err == context.Canceled {
		return tabletconn.Cancelled
	}
	return tabletconn.OperationalError(fmt.Sprintf("vttablet: %v", err))
}

func (sbc *sandboxConn) getError() error {
	if sbc.onConnUse != nil {
		sbc.onConnUse(sbc)
	}
	if sbc.mustFailRetry > 0 {
		sbc.mustFailRetry--
		return &tabletconn.ServerError{Code: tabletconn.ERR_RETRY, Err: "retry: err"}
	}
	if sbc.mustFailFatal > 0 {
		sbc.mustFailFatal--
		return &tabletconn.ServerError{Code: tabletconn.ERR_FATAL, Err: "fatal: err"}
	}
	if sbc.mustFailServer > 0 {
		sbc.mustFailServer--
		return &tabletconn.ServerError{Code: tabletconn.ERR_NORMAL, Err: "error: err"}
	}
	if sbc.mustFailConn > 0 {
		sbc.mustFailConn--
		return tabletconn.OperationalError(fmt.Sprintf("error: conn"))
	}
	if sbc.mustFailTxPool > 0 {
		sbc.mustFailTxPool--
		return &tabletconn.ServerError{Code: tabletconn.ERR_TX_POOL_FULL, Err: "tx_pool_full: err"}
	}
	if sbc.mustFailNotTx > 0 {
		sbc.mustFailNotTx--
		return &tabletconn.ServerError{Code: tabletconn.ERR_NOT_IN_TX, Err: "not_in_tx: err"}
	}
	return nil
}

// tabletErrorFromGRPC returns a tabletconn.ServerError or a
// tabletconn.OperationalError from the gRPC error.
func tabletErrorFromGRPC(err error) error {
	// TODO(aaijazi): Unfortunately, there's no better way to check for a gRPC server
	// error (vs a client error).
	// See: https://github.com/grpc/grpc-go/issues/319
	if !strings.Contains(err.Error(), vterrors.GRPCServerErrPrefix) {
		return tabletconn.OperationalError(fmt.Sprintf("vttablet: %v", err))
	}
	// server side error, convert it
	var code int
	switch grpc.Code(err) {
	case codes.Internal:
		code = tabletconn.ERR_FATAL
	case codes.FailedPrecondition:
		code = tabletconn.ERR_RETRY
	case codes.ResourceExhausted:
		code = tabletconn.ERR_TX_POOL_FULL
	case codes.Aborted:
		code = tabletconn.ERR_NOT_IN_TX
	default:
		code = tabletconn.ERR_NORMAL
	}

	return &tabletconn.ServerError{
		Code:       code,
		Err:        fmt.Sprintf("vttablet: %v", err),
		ServerCode: vterrors.GRPCCodeToErrorCode(grpc.Code(err)),
	}
}

func sandboxDialer(tablet *topodatapb.Tablet, timeout time.Duration) (tabletconn.TabletConn, error) {
	sand := getSandbox(tablet.Keyspace)
	sand.sandmu.Lock()
	defer sand.sandmu.Unlock()
	sand.DialCounter++
	if sand.DialMustFail > 0 {
		sand.DialMustFail--
		return nil, tabletconn.OperationalError(fmt.Sprintf("conn error"))
	}
	if sand.DialMustTimeout > 0 {
		time.Sleep(timeout)
		sand.DialMustTimeout--
		return nil, tabletconn.OperationalError(fmt.Sprintf("conn unreachable"))
	}
	sbc := sandboxconn.NewSandboxConn(tablet)
	return sbc, nil
}

// dialer is our tabletconn.Dialer
func dialer(ctx context.Context, tablet *topodatapb.Tablet, timeout time.Duration) (tabletconn.TabletConn, error) {
	t, ok := tabletMap[tablet.Alias.Uid]
	if !ok {
		return nil, tabletconn.OperationalError("connection refused")
	}

	return &internalTabletConn{
		tablet:     t,
		topoTablet: tablet,
	}, nil
}

func tabletErrorFromVitessError(ve *vterrors.VitessError) error {
	// see if the range is in the tablet error range
	if ve.Code >= vterrors.TabletError && ve.Code <= vterrors.UnknownTabletError {
		return &tabletconn.ServerError{
			Code: int(ve.Code - vterrors.TabletError),
			Err:  fmt.Sprintf("vttablet: %v", ve.Error()),
		}
	}

	return tabletconn.OperationalError(fmt.Sprintf("vttablet: %v", ve.Message))
}

// dialer is our tabletconn.Dialer
func dialer(ctx context.Context, endPoint *topodatapb.EndPoint, keyspace, shard string, tabletType topodatapb.TabletType, timeout time.Duration) (tabletconn.TabletConn, error) {
	tablet, ok := tabletMap[endPoint.Uid]
	if !ok {
		return nil, tabletconn.OperationalError("connection refused")
	}

	return &internalTabletConn{
		tablet:   tablet,
		endPoint: endPoint,
	}, nil
}

func sandboxDialer(ctx context.Context, endPoint topo.EndPoint, keyspace, shard string, timeout time.Duration) (tabletconn.TabletConn, error) {
	sand := getSandbox(keyspace)
	sand.sandmu.Lock()
	defer sand.sandmu.Unlock()
	sand.DialCounter++
	if sand.DialMustFail > 0 {
		sand.DialMustFail--
		return nil, tabletconn.OperationalError(fmt.Sprintf("conn error"))
	}
	if sand.DialMustTimeout > 0 {
		time.Sleep(timeout)
		sand.DialMustTimeout--
		return nil, tabletconn.OperationalError(fmt.Sprintf("conn unreachable"))
	}
	conns := sand.TestConns[shard]
	if conns == nil {
		panic(fmt.Sprintf("can't find shard %v", shard))
	}
	tconn := conns[endPoint.Uid]
	return tconn, nil
}

// tabletErrorFromRPCError reconstructs a tablet error from the
// RPCError, using the RPCError code.
func tabletErrorFromRPCError(rpcErr *pbv.RPCError) error {
	ve := vterrors.FromVtRPCError(rpcErr)

	// see if the range is in the tablet error range
	if ve.Code >= int64(pbv.ErrorCode_TabletError) && ve.Code <= int64(pbv.ErrorCode_UnknownTabletError) {
		return &tabletconn.ServerError{
			Code: int(ve.Code - int64(pbv.ErrorCode_TabletError)),
			Err:  fmt.Sprintf("vttablet: %v", ve.Error()),
		}
	}

	return tabletconn.OperationalError(fmt.Sprintf("vttablet: %v", ve.Message))
}

func (sbc *sandboxConn) getError() error {
	if sbc.onConnUse != nil {
		sbc.onConnUse(sbc)
	}
	if sbc.mustFailRetry > 0 {
		sbc.mustFailRetry--
		return &tabletconn.ServerError{
			Code:       tabletconn.ERR_RETRY,
			Err:        "retry: err",
			ServerCode: vtrpc.ErrorCode_QUERY_NOT_SERVED,
		}
	}
	if sbc.mustFailFatal > 0 {
		sbc.mustFailFatal--
		return &tabletconn.ServerError{
			Code:       tabletconn.ERR_FATAL,
			Err:        "fatal: err",
			ServerCode: vtrpc.ErrorCode_INTERNAL_ERROR,
		}
	}
	if sbc.mustFailServer > 0 {
		sbc.mustFailServer--
		return &tabletconn.ServerError{
			Code:       tabletconn.ERR_NORMAL,
			Err:        "error: err",
			ServerCode: vtrpc.ErrorCode_BAD_INPUT,
		}
	}
	if sbc.mustFailConn > 0 {
		sbc.mustFailConn--
		return tabletconn.OperationalError(fmt.Sprintf("error: conn"))
	}
	if sbc.mustFailTxPool > 0 {
		sbc.mustFailTxPool--
		return &tabletconn.ServerError{
			Code:       tabletconn.ERR_TX_POOL_FULL,
			Err:        "tx_pool_full: err",
			ServerCode: vtrpc.ErrorCode_RESOURCE_EXHAUSTED,
		}
	}
	if sbc.mustFailNotTx > 0 {
		sbc.mustFailNotTx--
		return &tabletconn.ServerError{
			Code:       tabletconn.ERR_NOT_IN_TX,
			Err:        "not_in_tx: err",
			ServerCode: vtrpc.ErrorCode_NOT_IN_TX,
		}
	}
	return nil
}

func sandboxDialer(context interface{}, endPoint topo.EndPoint, keyspace, shard string) (tabletconn.TabletConn, error) {
	sandmu.Lock()
	defer sandmu.Unlock()
	dialCounter++
	if dialMustFail > 0 {
		dialMustFail--
		return nil, tabletconn.OperationalError(fmt.Sprintf("conn error"))
	}
	tconn := testConns[endPoint.Uid]
	if tconn == nil {
		panic(fmt.Sprintf("can't find conn %v", endPoint.Uid))
	}
	tconn.(*sandboxConn).endPoint = endPoint
	return tconn, nil
}

func sandboxDialer(context context.Context, endPoint topo.EndPoint, keyspace, shard string, timeout time.Duration) (tabletconn.TabletConn, error) {
	sand := getSandbox(keyspace)
	sand.sandmu.Lock()
	defer sand.sandmu.Unlock()
	sand.DialCounter++
	if sand.DialMustFail > 0 {
		sand.DialMustFail--
		return nil, tabletconn.OperationalError(fmt.Sprintf("conn error"))
	}
	tconn := sand.TestConns[endPoint.Uid]
	if tconn == nil {
		panic(fmt.Sprintf("can't find conn %v", endPoint.Uid))
	}
	tconn.(*sandboxConn).endPoint = endPoint
	return tconn, nil
}

func (sbc *SandboxConn) getError() error {
	if sbc.MustFailRetry > 0 {
		sbc.MustFailRetry--
		return &tabletconn.ServerError{
			Err:        "retry: err",
			ServerCode: vtrpcpb.ErrorCode_QUERY_NOT_SERVED,
		}
	}
	if sbc.MustFailFatal > 0 {
		sbc.MustFailFatal--
		return &tabletconn.ServerError{
			Err:        "fatal: err",
			ServerCode: vtrpcpb.ErrorCode_INTERNAL_ERROR,
		}
	}
	if sbc.MustFailServer > 0 {
		sbc.MustFailServer--
		return &tabletconn.ServerError{
			Err:        "error: err",
			ServerCode: vtrpcpb.ErrorCode_BAD_INPUT,
		}
	}
	if sbc.MustFailConn > 0 {
		sbc.MustFailConn--
		return tabletconn.OperationalError(fmt.Sprintf("error: conn"))
	}
	if sbc.MustFailTxPool > 0 {
		sbc.MustFailTxPool--
		return &tabletconn.ServerError{
			Err:        "tx_pool_full: err",
			ServerCode: vtrpcpb.ErrorCode_RESOURCE_EXHAUSTED,
		}
	}
	if sbc.MustFailNotTx > 0 {
		sbc.MustFailNotTx--
		return &tabletconn.ServerError{
			Err:        "not_in_tx: err",
			ServerCode: vtrpcpb.ErrorCode_NOT_IN_TX,
		}
	}
	return nil
}

// withRetry sets up the connection and executes the action. If there are connection errors,
// it retries retryCount times before failing. It does not retry if the connection is in
// the middle of a transaction. While returning the error check if it maybe a result of
// a resharding event, and set the re-resolve bit and let the upper layers
// re-resolve and retry.
func (sdc *ShardConn) withRetry(ctx context.Context, action func(conn tabletconn.TabletConn) error, transactionID int64, isStreaming bool) error {
	var conn tabletconn.TabletConn
	var endPoint topo.EndPoint
	var err error
	var retry bool
	inTransaction := (transactionID != 0)
	// execute the action at least once even without retrying
	for i := 0; i < sdc.retryCount+1; i++ {
		conn, endPoint, err, retry = sdc.getConn(ctx)
		if err != nil {
			if retry {
				continue
			}
			return sdc.WrapError(err, endPoint, inTransaction)
		}
		// no timeout for streaming query
		if isStreaming {
			err = action(conn)
		} else {
			tmr := time.NewTimer(sdc.timeout)
			done := make(chan int)
			var errAction error
			go func() {
				errAction = action(conn)
				close(done)
			}()
			select {
			case <-tmr.C:
				err = tabletconn.OperationalError("vttablet: call timeout")
			case <-done:
				err = errAction
			}
			tmr.Stop()
		}
		if sdc.canRetry(err, transactionID, conn) {
			continue
		}
		return sdc.WrapError(err, endPoint, inTransaction)
	}
	return sdc.WrapError(err, endPoint, inTransaction)
}

// tabletErrorFromGRPC returns a tabletconn.OperationalError from the
// gRPC error.
func tabletErrorFromGRPC(err error) error {
	if grpc.Code(err) == codes.Internal {
		// server side error, convert it
		var code int
		errStr := err.Error()
		switch {
		case strings.Contains(errStr, "fatal: "):
			code = tabletconn.ERR_FATAL
		case strings.Contains(errStr, "retry: "):
			code = tabletconn.ERR_RETRY
		case strings.Contains(errStr, "tx_pool_full: "):
			code = tabletconn.ERR_TX_POOL_FULL
		case strings.Contains(errStr, "not_in_tx: "):
			code = tabletconn.ERR_NOT_IN_TX
		default:
			code = tabletconn.ERR_NORMAL
		}
		return &tabletconn.ServerError{Code: code, Err: fmt.Sprintf("vttablet: %v", err)}
	}
	return tabletconn.OperationalError(fmt.Sprintf("vttablet: %v", err))
}

func tabletError(err error) error {
	if err == nil {
		return nil
	}
	if _, ok := err.(rpcplus.ServerError); ok {
		var code int
		errStr := err.Error()
		switch {
		case strings.HasPrefix(errStr, "fatal"):
			code = tabletconn.ERR_FATAL
		case strings.HasPrefix(errStr, "retry"):
			code = tabletconn.ERR_RETRY
		case strings.HasPrefix(errStr, "tx_pool_full"):
			code = tabletconn.ERR_TX_POOL_FULL
		case strings.HasPrefix(errStr, "not_in_tx"):
			code = tabletconn.ERR_NOT_IN_TX
		default:
			code = tabletconn.ERR_NORMAL
		}
		return &tabletconn.ServerError{Code: code, Err: fmt.Sprintf("vttablet: %v", err)}
	}
	return tabletconn.OperationalError(fmt.Sprintf("vttablet: %v", err))
}

// tabletErrorFromGRPC returns a tabletconn.OperationalError from the
// gRPC error.
func tabletErrorFromGRPC(err error) error {
	return tabletconn.OperationalError(fmt.Sprintf("vttablet: %v", err))
}

func (sbc *SandboxConn) getError() error {
	if sbc.MustFailRetry > 0 {
		sbc.MustFailRetry--
		return &tabletconn.ServerError{
			Err:        "retry: err",
			ServerCode: vtrpcpb.ErrorCode_QUERY_NOT_SERVED,
		}
	}
	if sbc.MustFailFatal > 0 {
		sbc.MustFailFatal--
		return &tabletconn.ServerError{
			Err:        "fatal: err",
			ServerCode: vtrpcpb.ErrorCode_INTERNAL_ERROR,
		}
	}
	if sbc.MustFailServer > 0 {
		sbc.MustFailServer--
		return &tabletconn.ServerError{
			Err:        "error: err",
			ServerCode: vtrpcpb.ErrorCode_BAD_INPUT,
		}
	}
	if sbc.MustFailConn > 0 {
		sbc.MustFailConn--
		return tabletconn.OperationalError(fmt.Sprintf("error: conn"))
	}
	if sbc.MustFailTxPool > 0 {
		sbc.MustFailTxPool--
		return &tabletconn.ServerError{
			Err:        "tx_pool_full: err",
			ServerCode: vtrpcpb.ErrorCode_RESOURCE_EXHAUSTED,
		}
	}
	if sbc.MustFailNotTx > 0 {
		sbc.MustFailNotTx--
		return &tabletconn.ServerError{
			Err:        "not_in_tx: err",
			ServerCode: vtrpcpb.ErrorCode_NOT_IN_TX,
		}
	}
	if sbc.MustFailCanceled > 0 {
		sbc.MustFailCanceled--
		return &tabletconn.ServerError{
			Err:        "canceled: err",
			ServerCode: vtrpcpb.ErrorCode_CANCELLED,
		}
	}
	if sbc.MustFailUnknownError > 0 {
		sbc.MustFailUnknownError--
		return &tabletconn.ServerError{
			Err:        "unknown error: err",
			ServerCode: vtrpcpb.ErrorCode_UNKNOWN_ERROR,
		}
	}
	if sbc.MustFailDeadlineExceeded > 0 {
		sbc.MustFailDeadlineExceeded--
		return &tabletconn.ServerError{
			Err:        "deadline exceeded: err",
			ServerCode: vtrpcpb.ErrorCode_DEADLINE_EXCEEDED,
		}
	}
	if sbc.MustFailIntegrityError > 0 {
		sbc.MustFailIntegrityError--
		return &tabletconn.ServerError{
			Err:        "integrity error: err",
			ServerCode: vtrpcpb.ErrorCode_INTEGRITY_ERROR,
		}
	}
	if sbc.MustFailPermissionDenied > 0 {
		sbc.MustFailPermissionDenied--
		return &tabletconn.ServerError{
			Err:        "permission denied: err",
			ServerCode: vtrpcpb.ErrorCode_PERMISSION_DENIED,
		}
	}
	if sbc.MustFailTransientError > 0 {
		sbc.MustFailTransientError--
		return &tabletconn.ServerError{
			Err:        "transient error: err",
			ServerCode: vtrpcpb.ErrorCode_TRANSIENT_ERROR,
		}
	}
	if sbc.MustFailUnauthenticated > 0 {
		sbc.MustFailUnauthenticated--
		return &tabletconn.ServerError{
			Err:        "unauthenticated: err",
			ServerCode: vtrpcpb.ErrorCode_UNAUTHENTICATED,
		}
	}

	return nil
}