#!/bin/bash

# Copyright 2024 TiKV Authors
#
# 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.

output="cmds_generated.go"

cat <<EOF > $output
// Code generated by gen.sh. DO NOT EDIT.

package tikvrpc

import (
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
)
EOF

cmds=(
  Get
  Scan
  Prewrite
  PessimisticLock
  PessimisticRollback
  Commit
  Cleanup
  BatchGet
  BatchRollback
  ScanLock
  ResolveLock
  GC
  DeleteRange
  RawGet
  RawBatchGet
  RawPut
  RawBatchPut
  RawDelete
  RawBatchDelete
  RawDeleteRange
  RawScan
  RawGetKeyTTL
  RawCompareAndSwap
  RawChecksum
  UnsafeDestroyRange
  RegisterLockObserver
  CheckLockObserver
  RemoveLockObserver
  PhysicalScanLock
  Cop
  BatchCop
  MvccGetByKey
  MvccGetByStartTs
  SplitRegion
  TxnHeartBeat
  CheckTxnStatus
  CheckSecondaryLocks
  FlashbackToVersion
  PrepareFlashbackToVersion
  Flush
  BufferBatchGet
)

cat <<EOF >> $output

func patchCmdCtx(req *Request, cmd CmdType, ctx *kvrpcpb.Context) bool {
	switch cmd {
EOF

for cmd in "${cmds[@]}"; do
cat <<EOF >> $output
	case Cmd${cmd}:
		if req.rev == 0 {
			req.${cmd}().Context = ctx
		} else {
			cmd := *req.${cmd}()
			cmd.Context = ctx
			req.Req = &cmd
		}
		req.rev++
EOF
done

cat <<EOF >> $output
	default:
		return false
	}
	return true
}
EOF