// Copyright 2023 PingCAP, Inc.
//
// 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 remote

import (
	"context"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cockroachdb/pebble"
	"github.com/google/uuid"
	"github.com/pingcap/errors"
	"github.com/pingcap/tidb/br/pkg/pdutil"
	"github.com/pingcap/tidb/pkg/lightning/backend"
	"github.com/pingcap/tidb/pkg/lightning/backend/encode"
	"github.com/pingcap/tidb/pkg/lightning/backend/kv"
	"github.com/pingcap/tidb/pkg/lightning/backend/local"
	"github.com/pingcap/tidb/pkg/lightning/checkpoints"
	"github.com/pingcap/tidb/pkg/lightning/common"
	"github.com/pingcap/tidb/pkg/lightning/config"
	"github.com/pingcap/tidb/pkg/lightning/errormanager"
	"github.com/pingcap/tidb/pkg/lightning/log"
	"github.com/pingcap/tidb/pkg/lightning/metric"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tikv/client-go/v2/oracle"
	tikvclient "github.com/tikv/client-go/v2/tikv"
	pd "github.com/tikv/pd/client"
	"go.uber.org/zap"
)

/*
	Remote load data worker API:

	1. init task:
		POST /load_data?cluster_id=%d&task_id=%s&start_ts=%d&commit_ts=%d

	2. put chunk:
		PUT /load_data?cluster_id=%d&task_id=%s&writer_id=%d&chunk_id=%d

		key_len(2) + key(key_len) + val_len(4) + value(val_len) + row_id_len(2) + row_id(row_id_len)
		key_len(2) + key(key_len) + val_len(4) + value(val_len) + row_id_len(2) + row_id(row_id_len)
		...

	3. flush:
		POST /load_data?cluster_id=%d&task_id=%s&flush=true

	4. build task:
		POST /load_data?cluster_id=%d&task_id=%s&build=true&compression=zstd&split_size=%d&split_keys=%d

	5. get task states:
		GET /load_data?cluster_id=%d&task_id=%s

		{"canceled": false, "finished": false, "error": "", "created-files": 10, "ingested-regions": 3}

	6. clean up task:
		DELETE /load_data?cluster_id=%d&task_id=%s
*/

const (
	duplicateDBName = "duplicates"
	maxConnsPerHost = 10
)

// LoadDataStates is json data that returned by remote server GET API.
type LoadDataStates struct {
	TaskID          string `json:"task-id"`
	Canceled        bool   `json:"canceled"`
	Finished        bool   `json:"finished"`
	Error           string `json:"error"`
	CreatedFiles    int    `json:"created-files"`
	IngestedRegions int    `json:"ingested-regions"`
	TotalKVs        int64  `json:"total-kvs"`

	DuplicateEntries []duplicateEntry `json:"duplicated-entries"`
}

type duplicateEntry struct {
	Key    string   `json:"key"`
	Values []string `json:"values"`
}

func (s *LoadDataStates) hasDuplicateEntries() bool {
	return len(s.DuplicateEntries) > 0
}

// NewBackend creates a new remote backend instance.
func NewBackend(
	ctx context.Context,
	tls *common.TLS,
	cfg *BackendConfig,
	pdSvcDiscovery pd.ServiceDiscovery,
) (backend.Backend, error) {
	localFile := cfg.SortedKVDir

	shouldCreate := true
	if cfg.CheckpointEnabled {
		if info, err := os.Stat(localFile); err != nil {
			if !os.IsNotExist(err) {
				return nil, err
			}
		} else if info.IsDir() {
			shouldCreate = false
		}
	}
	if shouldCreate {
		err := os.Mkdir(localFile, 0o700)
		if err != nil {
			return nil, common.ErrInvalidSortedKVDir.Wrap(err).GenWithStackByArgs(localFile)
		}
	}
	var (
		duplicateDB *pebble.DB
		err         error
	)
	keyAdapter := common.KeyAdapter(common.NoopKeyAdapter{})
	duplicateDetection := cfg.DuplicateResolution != config.NoneOnDup
	if duplicateDetection {
		duplicateDB, err = local.OpenDuplicateDB(localFile)
		if err != nil {
			return nil, common.ErrOpenDuplicateDB.Wrap(err).GenWithStackByArgs()
		}
		keyAdapter = common.DupDetectKeyAdapter{}
	}
	reportErrOnDup := cfg.DuplicateResolution == config.ErrorOnDup

	var pdAddrs []string
	if pdSvcDiscovery != nil {
		pdAddrs = pdSvcDiscovery.GetServiceURLs()
	} else {
		pdAddrs = strings.Split(cfg.PdAddr, ",")
	}
	pdCtl, err := pdutil.NewPdController(ctx, cfg.KeyspaceName, pdAddrs, tls.TLSConfig(), tls.ToPDSecurityOption())
	if err != nil {
		log.L().Error("fail to create pd controller", zap.Error(err))
		return nil, common.ErrCreatePDClient.Wrap(err).GenWithStackByArgs()
	}

	var pdCliForTiKV *tikvclient.CodecPDClient
	if len(cfg.KeyspaceName) == 0 {
		pdCliForTiKV = tikvclient.NewCodecPDClient(tikvclient.ModeTxn, pdCtl.GetPDClient())
	} else {
		pdCliForTiKV, err = tikvclient.NewCodecPDClientWithKeyspace(tikvclient.ModeTxn, pdCtl.GetPDClient(), cfg.KeyspaceName)
		if err != nil {
			log.L().Error("fail to create pd cli", zap.Error(err))
			return nil, common.ErrCreatePDClient.Wrap(err).GenWithStackByArgs()
		}
	}

	tikvCodec := pdCliForTiKV.GetCodec()
	spkv, err := tikvclient.NewEtcdSafePointKV(strings.Split(cfg.PdAddr, ","), tls.TLSConfig())
	if err != nil {
		log.L().Error("fail to create etcd client", zap.Error(err))
		return nil, common.ErrCreateKVClient.Wrap(err).GenWithStackByArgs()
	}
	rpcCli := tikvclient.NewRPCClient(tikvclient.WithSecurity(tls.ToTiKVSecurityConfig()), tikvclient.WithCodec(tikvCodec))
	tikvCli, err := tikvclient.NewKVStore("lightning-remote-backend", pdCliForTiKV, spkv, rpcCli)
	if err != nil {
		log.L().Error("fail to create TiKV client", zap.Error(err))
		return nil, common.ErrCreateKVClient.Wrap(err).GenWithStackByArgs()
	}

	httpClient := http.DefaultClient
	if tlsConfig := tls.TLSConfig(); tlsConfig != nil {
		httpClient = &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: tlsConfig,
			},
		}
	}

	ks := pdCliForTiKV.GetCodec().GetKeyspace()
	ksID := pdCliForTiKV.GetCodec().GetKeyspaceID()
	remote := &Backend{
		workerAddr: cfg.TikvImporterAddr,
		pdCtl:      pdCtl,
		tls:        tls,
		pdAddr:     cfg.PdAddr,
		keyspace:   ks,
		keyspaceID: uint32(ksID),
		logger:     log.FromContext(ctx),
		httpClient: httpClient,

		tikvCodec:          tikvCodec,
		tikvCli:            tikvCli,
		duplicateDB:        duplicateDB,
		keyAdapter:         keyAdapter,
		duplicateDetection: duplicateDetection,
		reportErrOnDup:     reportErrOnDup,
		checkpointEnabled:  cfg.CheckpointEnabled,
		localStoreDir:      localFile,
		chunkCacheDir:      cfg.ChunkCacheDir,
		chunkCacheInMem:    cfg.ChunkCacheInMem,
	}

	if m, ok := metric.FromContext(ctx); ok {
		remote.metrics = m
	}

	return remote, nil
}

// BackendConfig is the config for remote backend.
type BackendConfig struct {
	PdAddr              string
	TikvImporterAddr    string
	KeyspaceName        string
	SortedKVDir         string
	ChunkCacheDir       string
	ChunkCacheInMem     bool
	CheckpointEnabled   bool
	WorkerConcurrency   int
	DuplicateResolution config.DuplicateResolutionAlgorithm
}

// NewBackendConfig creates a new BackendConfig.
func NewBackendConfig(cfg *config.Config, keyspaceName string) *BackendConfig {
	return &BackendConfig{
		PdAddr:              cfg.TiDB.PdAddr,
		TikvImporterAddr:    cfg.TikvImporter.Addr,
		KeyspaceName:        keyspaceName,
		SortedKVDir:         cfg.TikvImporter.SortedKVDir,
		ChunkCacheDir:       cfg.TikvImporter.ChunkCacheDir,
		CheckpointEnabled:   cfg.Checkpoint.Enable,
		WorkerConcurrency:   cfg.TikvImporter.RangeConcurrency * 2,
		DuplicateResolution: cfg.Conflict.Strategy,
	}
}

// Backend is a remote backend that sends KV pairs to remote worker.
type Backend struct {
	workerAddr string
	pdCtl      *pdutil.PdController
	tls        *common.TLS
	pdAddr     string
	keyspace   []byte
	keyspaceID uint32
	metrics    *metric.Metrics
	engines    sync.Map
	logger     log.Logger
	httpClient *http.Client

	tikvCodec          tikvclient.Codec
	duplicateDB        *pebble.DB
	tikvCli            *tikvclient.KVStore
	keyAdapter         common.KeyAdapter
	duplicateDetection bool
	reportErrOnDup     bool
	checkpointEnabled  bool
	localStoreDir      string
	chunkCacheDir      string
	chunkCacheInMem    bool
}

// Close the connection to the backend.
func (b *Backend) Close() {
	if b.duplicateDB != nil {
		// Check if there are duplicates that are not collected.
		iter, err := b.duplicateDB.NewIter(&pebble.IterOptions{})
		if err == nil {
			hasDuplicates := iter.First()
			allIsWell := true
			if err := iter.Error(); err != nil {
				b.logger.Warn("iterate duplicate db failed", zap.Error(err))
				allIsWell = false
			}
			if err := iter.Close(); err != nil {
				b.logger.Warn("close duplicate db iter failed", zap.Error(err))
				allIsWell = false
			}
			if err := b.duplicateDB.Close(); err != nil {
				b.logger.Warn("close duplicate db failed", zap.Error(err))
				allIsWell = false
			}
			// If checkpoint is disabled, or we don't detect any duplicate, then this duplicate
			// db dir will be useless, so we clean up this dir.
			if allIsWell && (!b.checkpointEnabled || !hasDuplicates) {
				if err := os.RemoveAll(filepath.Join(b.localStoreDir, duplicateDBName)); err != nil {
					b.logger.Warn("remove duplicate db file failed", zap.Error(err))
				}
			}
			b.duplicateDB = nil
		} else {
			b.logger.Warn("create iter failed", zap.Error(err))
		}
	}

	// if checkpoint is disable or we finish load all data successfully, then files in this
	// dir will be useless, so we clean up this dir and all files in it.
	if !b.checkpointEnabled || common.IsEmptyDir(b.localStoreDir) {
		err := os.RemoveAll(b.localStoreDir)
		if err != nil {
			b.logger.Warn("remove local db file failed", zap.Error(err))
		}
	}
	_ = b.tikvCli.Close()
	b.pdCtl.Close()
}

// RetryImportDelay returns the duration to sleep when retrying an import
func (b *Backend) RetryImportDelay() time.Duration {
	return 0
}

// ShouldPostProcess returns whether KV-specific post-processing should be
// performed for this backend. Post-processing includes checksum and analyze.
func (b *Backend) ShouldPostProcess() bool {
	return true
}

// OpenEngine opens an engine for writing.
func (b *Backend) OpenEngine(ctx context.Context, cfg *backend.EngineConfig, engineUUID uuid.UUID) error {
	physical, logical, err := b.pdCtl.GetPDClient().GetTS(ctx)
	if err != nil {
		return err
	}
	ts := oracle.ComposeTS(physical, logical)

	loadDataTaskID := genLoadDataTaskID(b.keyspaceID, cfg)
	var retryCounter prometheus.Counter
	if b.metrics != nil {
		retryCounter = b.metrics.RemoteClientRetryCounter.WithLabelValues(loadDataTaskID)
	}
	e, _ := b.engines.LoadOrStore(engineUUID, &engine{
		ctx:            ctx,
		logger:         log.With(zap.String("loadDataTaskID", loadDataTaskID)),
		loadDataTaskID: loadDataTaskID,
		ts:             ts,
		tbl:            cfg.TableInfo,
		addr:           b.workerAddr,
		clusterID:      b.pdCtl.GetPDClient().GetClusterID(ctx),
		writers:        sync.Map{},
		httpClient:     b.httpClient,
		backend:        b,

		retryCounter: retryCounter,
	})
	engine := e.(*engine)
	if cfg.ValidCheckpoint {
		exist, err := b.checkLoadDataTask(ctx, cfg)
		if err != nil {
			return err
		}
		if !exist {
			b.logger.Error("task not found in remote", zap.String("loadDataTaskID", loadDataTaskID))
			return ErrTaskNotFound
		}
	}

	if engine.loadDataTaskID == loadDataTaskID {
		// newly created engine.
		err = b.loadDataInit(ctx, engine, cfg.EstimatedDataSize)
		if err != nil {
			b.engines.Delete(engineUUID)
			return err
		}
	}
	return nil
}

func (b *Backend) loadDataInit(ctx context.Context, engine *engine, dataSize int64) error {
	for {
		url := fmt.Sprintf(
			"%s/load_data?cluster_id=%d&task_id=%s&start_ts=%d&commit_ts=%d&data_size=%d&new_client=true",
			engine.addr,
			engine.clusterID,
			engine.loadDataTaskID,
			engine.ts,
			engine.ts+1,
			dataSize,
		)
		client := *b.httpClient
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		}

		resp, err := sendRequestWithRetry(ctx, &client, "POST", url, nil, engine.retryCounter)
		if err != nil {
			return err
		}
		defer resp.Body.Close()

		if resp.StatusCode == http.StatusFound {
			engine.addr = strings.TrimSuffix(resp.Header.Get("Location"), "/load_data")
			b.logger.Info("redirect to loadData worker",
				zap.String("loadDataTaskID", engine.loadDataTaskID),
				zap.String("worker addr", engine.addr))
			continue
		}
		if resp.StatusCode != http.StatusOK {
			msg, err := io.ReadAll(resp.Body)
			// If the task exists, we can continue to import data.
			if err == nil {
				if strings.TrimSpace(string(msg)) == taskExitsMsg {
					b.logger.Info("loadData task has inited in load_data worker",
						zap.String("loadDataTask", engine.loadDataTaskID),
						zap.String("worker addr", engine.addr))
				}
				// msg format as `data size %d exceeds max data size %d`
				if strings.Contains(string(msg), taskExceedMaxSizeMsg) {
					return errors.New(string(msg) + ". Please contact support to raise threshold.")
				}
				return nil
			}
			b.logger.Error("failed to init loadData task",
				zap.String("loadDataTaskID", engine.loadDataTaskID),
				zap.String("status", resp.Status),
				zap.String("msg", string(msg)))
			return ErrInitTask
		}
		return nil
	}
}

// CloseEngine closes backend engine by uuid.
func (b *Backend) CloseEngine(ctx context.Context, cfg *backend.EngineConfig, engineUUID uuid.UUID) error {
	engine, err := b.getEngine(engineUUID)
	if err != nil {
		exist, err := b.checkLoadDataTask(ctx, cfg)
		if err != nil {
			return err
		}
		if !exist {
			b.logger.Error("task not found in remote", zap.String("loadDataTaskID", genLoadDataTaskID(b.keyspaceID, cfg)))
			return ErrTaskNotFound
		}
		b.logger.Info("not found engine in local, open a new one")
		err = b.OpenEngine(ctx, cfg, engineUUID)
		if err != nil {
			return err
		}
		engine, err = b.getEngine(engineUUID)
		if err != nil {
			b.logger.Warn("failed to open a new engine", zap.Error(err))
			return err
		}
	}

	err = engine.flush(ctx)
	if err != nil {
		return err
	}

	engine.writers.Range(func(key, value any) bool {
		sender := key.(*writer).chunkSender
		err = sender.close()
		if err != nil {
			b.logger.Warn("failed to close chunk sender", zap.Error(err), zap.Uint64("sender", sender.id))
		}
		return true
	})

	if err == nil {
		engine.writers.Clear()
	}

	return err
}

func (b *Backend) checkLoadDataTask(ctx context.Context, cfg *backend.EngineConfig) (bool, error) {
	clusterID := b.pdCtl.GetPDClient().GetClusterID(ctx)
	loadDataTaskID := genLoadDataTaskID(b.keyspaceID, cfg)
	addr := b.workerAddr

	var retryCounter prometheus.Counter
	if b.metrics != nil {
		retryCounter = b.metrics.RemoteClientRetryCounter.WithLabelValues(loadDataTaskID)
	}
	for {
		url := fmt.Sprintf("%s/load_data?cluster_id=%d&task_id=%s", addr, clusterID, loadDataTaskID)
		client := *b.httpClient
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		}

		resp, err := sendRequestWithRetry(ctx, &client, "GET", url, nil, retryCounter)
		if err != nil {
			return false, err
		}
		defer resp.Body.Close()

		if resp.StatusCode == http.StatusFound {
			addr = strings.TrimSuffix(resp.Header.Get("Location"), "/load_data")
			b.logger.Info("redirect to loadData worker",
				zap.String("loadDataTaskID", loadDataTaskID),
				zap.String("worker addr", addr))
			continue
		}
		if resp.StatusCode != http.StatusOK {
			if resp.StatusCode == http.StatusNotFound {
				return false, nil
			}
			msg, err := io.ReadAll(resp.Body)
			if err != nil {
				return false, err
			}
			b.logger.Warn("failed to get task", zap.String("status", resp.Status), zap.String("msg", string(msg)))
			return false, ErrGetTaskState
		}
		return true, nil
	}
}

// ImportEngine imports an engine to TiKV.
func (b *Backend) ImportEngine(ctx context.Context, engineUUID uuid.UUID, regionSplitSize, regionSplitKeys int64) error {
	engine, err := b.getEngine(engineUUID)
	if err != nil {
		return err
	}
	err = b.loadDataBuild(ctx, engine, regionSplitSize, regionSplitKeys)
	if err != nil {
		return err
	}
	start := time.Now()
	ticker := time.NewTicker(time.Second * 1)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			states, err := b.loadDataGetStates(ctx, engine)
			if err != nil {
				return err
			}
			if states.Canceled {
				b.logger.Error("loadData canceled",
					zap.String("loadDataTaskID", engine.loadDataTaskID),
					zap.String("error", states.Error))
				if strings.Contains(states.Error, "Ingest::Overlap: region has overlap data") {
					return ErrOverlap
				} else if strings.Contains(states.Error, "TooManyDuplicatedKeys") {
					return ErrDuplicateKey
				}
				return ErrTaskCanceled
			} else if states.Finished {
				if states.hasDuplicateEntries() {
					err = b.handleDuplicateEntries(ctx, engine, states)
					if err != nil {
						return err
					}
				}
				engine.importedKVs.Store(states.TotalKVs)
				b.logger.Info("loadData finished",
					zap.String("db", engine.tbl.DB),
					zap.String("table", engine.tbl.Name),
					zap.String("loadDataTaskID", engine.loadDataTaskID),
					zap.Int64("imported kvs", states.TotalKVs),
					zap.Duration("elapsed", time.Since(start)))
				return nil
			} else {
				b.logger.Info(
					"loadData states",
					zap.String("db", engine.tbl.DB),
					zap.String("table", engine.tbl.Name),
					zap.String("loadDataTaskID", engine.loadDataTaskID),
					zap.Int64("total kvs", states.TotalKVs),
					zap.Int("created files", states.CreatedFiles),
					zap.Int("ingested regions", states.IngestedRegions))
			}
		}
	}
}

// GetImportedKVCount returns the number of imported KV pairs of some engine.
func (b *Backend) GetImportedKVCount(engineUUID uuid.UUID) int64 {
	v, ok := b.engines.Load(engineUUID)
	if !ok {
		return 0
	}
	e := v.(*engine)
	return e.importedKVs.Load()
}

func (b *Backend) handleDuplicateEntries(_ context.Context, engine *engine, states *LoadDataStates) error {
	if !b.duplicateDetection {
		return nil
	}

	b.logger.Info("handling duplicate entries",
		zap.String("db", engine.tbl.DB),
		zap.String("table", engine.tbl.Name),
		zap.Int("duplicateEntries", len(states.DuplicateEntries)))

	if b.reportErrOnDup {
		dupKey, err := hex.DecodeString(states.DuplicateEntries[0].Key)
		dupKey, err = b.tikvCodec.DecodeKey(dupKey)
		if err != nil {
			b.logger.Warn("failed to decode key", zap.String("key", states.DuplicateEntries[0].Key), zap.Error(err))
		}

		dupVal, err := hex.DecodeString(states.DuplicateEntries[0].Values[0])
		if err != nil {
			b.logger.Warn("failed to decode value", zap.String("value", states.DuplicateEntries[0].Values[0]), zap.Error(err))
		}

		return common.ErrFoundDuplicateKeys.FastGenByArgs(dupKey, dupVal)
	}

	writeBatch := b.duplicateDB.NewBatch()
	writeBatchSize := int64(0)
	for _, entry := range states.DuplicateEntries {
		if len(entry.Values) == 0 {
			continue
		}
		rawKey, err := hex.DecodeString(entry.Key)
		if err != nil {
			b.logger.Warn("failed to decode key", zap.String("key", entry.Key), zap.Error(err))
			return err
		}
		for i, value := range entry.Values {
			// append index to rawKey to make it unique.
			encodedRawKey := b.keyAdapter.Encode(nil, rawKey, common.EncodeIntRowID(int64(i)))

			rawValue, err := hex.DecodeString(value)
			if err != nil {
				b.logger.Warn("failed to decode value", zap.String("value", value), zap.Error(err))
				return err
			}
			writeBatch.Set(encodedRawKey, rawValue, nil)
			writeBatchSize += int64(len(encodedRawKey) + len(rawValue))
			if writeBatchSize >= maxDuplicateBatchSize {
				err = writeBatch.Commit(nil)
				if err != nil {
					b.logger.Warn("failed to write duplicate entries", zap.Error(err))
					return err
				}
				writeBatch = b.duplicateDB.NewBatch()
				writeBatchSize = 0
			}
		}
	}

	if writeBatchSize > 0 {
		err := writeBatch.Commit(nil)
		if err != nil {
			b.logger.Warn("failed to write duplicate entries", zap.Error(err))
			return err
		}
	}
	return nil
}

func (b *Backend) allocateSenderID(ctx context.Context) (uint64, error) {
	physical, logical, err := b.pdCtl.GetPDClient().GetTS(ctx)
	if err != nil {
		return 0, err
	}
	ts := oracle.ComposeTS(physical, logical)
	return ts, nil
}

func (b *Backend) loadDataBuild(ctx context.Context, engine *engine, splitSize, splitKeys int64) error {
	url := fmt.Sprintf("%s/load_data?cluster_id=%d&task_id=%s&build=true&compression=zstd&split_size=%d&split_keys=%d",
		engine.addr, engine.clusterID, engine.loadDataTaskID, splitSize, splitKeys)

	_, err := sendRequest(ctx, b.httpClient, "POST", url, nil, engine.retryCounter)
	return err
}

func (b *Backend) loadDataGetStates(ctx context.Context, engine *engine) (*LoadDataStates, error) {
	url := fmt.Sprintf("%s/load_data?cluster_id=%d&task_id=%s", engine.addr, engine.clusterID, engine.loadDataTaskID)
	data, err := sendRequest(ctx, b.httpClient, "GET", url, nil, engine.retryCounter)
	if err != nil {
		return nil, err
	}
	state := new(LoadDataStates)
	err = json.Unmarshal(data, state)
	if err != nil {
		return nil, err
	}
	return state, nil
}

// CleanupEngine cleanup the engine and reclaim the space.
func (b *Backend) CleanupEngine(ctx context.Context, engineUUID uuid.UUID) error {
	engine, err := b.getEngine(engineUUID)
	if err != nil {
		return err
	}
	return b.loadDataCleanUp(ctx, engine)
}

func (b *Backend) loadDataCleanUp(ctx context.Context, engine *engine) error {
	url := fmt.Sprintf("%s/load_data?cluster_id=%d&task_id=%s", engine.addr, engine.clusterID, engine.loadDataTaskID)
	_, err := sendRequest(ctx, b.httpClient, "DELETE", url, nil, engine.retryCounter)
	return err
}

// FlushEngine ensures all KV pairs written to an open engine has been
// synchronized, such that kill-9'ing Lightning afterwards and resuming from
// checkpoint can recover the exact same content.
//
// This method is only relevant for local backend, and is no-op for all
// other backends.
func (b *Backend) FlushEngine(ctx context.Context, engineUUID uuid.UUID) error {
	engine, err := b.getEngine(engineUUID)
	if err != nil {
		return err
	}
	return engine.flush(ctx)
}

// FlushAllEngines performs FlushEngine on all opened engines. This is a
// very expensive operation and should only be used in some rare situation
// (e.g. preparing to resolve a disk quota violation).
// This method is only relevant for local backend, and is no-op for remote backend.
func (b *Backend) FlushAllEngines(ctx context.Context) error {
	return nil
}

// EngineFileSizes obtains the size occupied locally of all engines managed
// by this backend. This method is used to compute disk quota.
// It can return nil if the content are all stored remotely.
func (b *Backend) EngineFileSizes() []backend.EngineFileSize {
	return nil
}

// ResetEngine clears all written KV pairs in this opened engine.
func (b *Backend) ResetEngine(ctx context.Context, engineUUID uuid.UUID) error {
	return errors.New("cannot reset an engine in Remote backend")
}

// LocalWriter obtains a thread-local EngineWriter for writing rows into the given engine.
func (b *Backend) LocalWriter(ctx context.Context, cfg *backend.LocalWriterConfig, engineUUID uuid.UUID) (backend.EngineWriter, error) {
	engine, err := b.getEngine(engineUUID)
	if err != nil {
		return nil, err
	}
	// If lightning pod is restarted, there is no guarantee that the chunk generated is same as before, so we need to
	// allocate a new sender id instead of using writer's id.
	senderID, err := b.allocateSenderID(ctx)
	if err != nil {
		return nil, err
	}
	cache, err := newChunkCache(engine.loadDataTaskID, senderID, b.chunkCacheDir, b.chunkCacheInMem)
	if err != nil {
		return nil, err
	}
	writer := &writer{
		e:           engine,
		keyspace:    b.keyspace,
		chunkSender: newChunkSender(ctx, senderID, engine, cache),
	}
	engine.writers.Store(writer, struct{}{})
	return writer, nil
}

func (b *Backend) getEngine(engineUUID uuid.UUID) (*engine, error) {
	v, ok := b.engines.Load(engineUUID)
	if !ok {
		return nil, errors.Errorf("could not found engine for %s", engineUUID.String())
	}
	return v.(*engine), nil
}

type engine struct {
	ctx            context.Context
	logger         log.Logger
	loadDataTaskID string
	ts             uint64
	tbl            *checkpoints.TidbTableInfo
	writers        sync.Map
	addr           string
	clusterID      uint64

	httpClient   *http.Client
	backend      *Backend
	retryCounter prometheus.Counter
	importedKVs  atomic.Int64
}

func (e *engine) flush(ctx context.Context) error {
	var err error
	e.writers.Range(func(key, _ any) bool {
		w := key.(*writer)
		err = w.addChunk(ctx)
		if err != nil {
			return false
		}

		err = w.chunkSender.flush(ctx)
		if err != nil {
			return false
		}
		return true
	})
	if err != nil {
		e.logger.Warn("failed to flush writer", zap.Error(err))
		return errors.Trace(err)
	}

	return err
}

// writer implments the EngineWriter interface
type writer struct {
	e           *engine
	chunkSender *chunkSender
	lastChunkID uint64

	buf      []byte
	keyspace []byte
}

const batchSize int = 8 * 1024 * 1024

func (w *writer) AppendRows(
	ctx context.Context,
	columnNames []string,
	rows encode.Rows,
) error {
	kvs := kv.Rows2KvPairs(rows)
	if len(kvs) == 0 {
		return nil
	}
	lenBuf := make([]byte, 4)
	for _, pair := range kvs {
		binary.LittleEndian.PutUint16(lenBuf[:2], uint16(len(pair.Key)+len(w.keyspace)))
		w.buf = append(w.buf, lenBuf[:2]...)
		w.buf = append(w.buf, w.keyspace...)
		w.buf = append(w.buf, pair.Key...)
		binary.LittleEndian.PutUint32(lenBuf, uint32(len(pair.Val)))
		w.buf = append(w.buf, lenBuf...)
		w.buf = append(w.buf, pair.Val...)
		binary.LittleEndian.PutUint16(lenBuf, uint16(len(pair.RowID)))
		w.buf = append(w.buf, lenBuf[:2]...)
		w.buf = append(w.buf, pair.RowID...)
	}
	if len(w.buf) > batchSize {
		return w.addChunk(ctx)
	}
	return nil
}

func (w *writer) addChunk(ctx context.Context) error {
	if len(w.buf) == 0 {
		// nothing to do
		return nil
	}

	chunkID, err := w.chunkSender.putChunk(ctx, w.buf)
	if err != nil {
		return err
	}
	w.lastChunkID = chunkID
	w.buf = w.buf[:0]

	return nil
}

func (w *writer) IsSynced() bool {
	return len(w.buf) == 0 && w.lastChunkID <= w.chunkSender.getFlushedChunkID()
}

func (w *writer) Close(ctx context.Context) (backend.ChunkFlushStatus, error) {
	if len(w.buf) > 0 {
		err := w.addChunk(ctx)
		if err != nil {
			return w, err
		}
	}
	w.buf = nil
	return w, nil
}

func (w *writer) Flushed() bool {
	return w.lastChunkID <= w.chunkSender.getFlushedChunkID()
}

// GetDupeController returns a new dupe controller.
func (b *Backend) GetDupeController(dupeConcurrency int, errorMgr *errormanager.ErrorManager) *local.DupeController {
	return local.NewDupeController(dupeConcurrency, errorMgr, nil, b.tikvCli, b.tikvCodec, b.duplicateDB, b.keyAdapter, nil, false)
}