package handle

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/kvproto/pkg/keyspacepb"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/kvproto/pkg/metapb"
	"github.com/pingcap/log"
	"github.com/pingcap/tidb/pkg/util/codec"
	"github.com/tidbcloud/aws-shared-provider/pkg/datacleaner/common"
	"github.com/tidbcloud/aws-shared-provider/pkg/datacleaner/placement"
	pdhttp "github.com/tidbcloud/aws-shared-provider/pkg/utils/pd"
	"github.com/tikv/client-go/v2/metrics"
	"github.com/tikv/client-go/v2/tikvrpc"
	"github.com/tikv/client-go/v2/txnkv"
	pd "github.com/tikv/pd/client"
	"go.uber.org/zap"
)

const (
	unsafeDestroyRangeTimeout = 5 * time.Minute
	keyspacePrefixLen         = 4
)

var (
	txnKeyspaceStartBoundary = []byte{'x', 0, 0, 0}
)

func DeleteRange(ctx context.Context, pdClient pd.Client, client *txnkv.Client,
	record *common.KeyspaceRecord, isRun, isSkipConfirm bool) (successed bool, err error) {
	keyspaceId := record.KeyspaceId
	log.Info("[BEGIN] UnsafeDestroyRange", zap.Uint32("keyspaceId", keyspaceId))
	// deleteRanges
	rawLeftBound, rawRightBound, txnLeftBound, txnRightBound := common.GetRange(keyspaceId)

	var confirmMsg string
	if isSkipConfirm {
		confirmMsg = "yes"
	} else {
		fmt.Println("Please confirm is't needs to be GC.(yes/no)")
		_, _ = fmt.Scanln(&confirmMsg)
	}

	if confirmMsg == "yes" {
		if isRun {
			log.Info("[REAL RUN]", zap.Uint32("keyspaceId", keyspaceId))
			// clean rawkv range
			err = UnsafeDestroyRange(ctx, pdClient, rawLeftBound, rawRightBound, client)
			if err != nil {
				log.Error("UnsafeDestroyRange raw mode range failed", zap.Error(err))
				return successed, err
			}
			// clean txnkv range
			err = UnsafeDestroyRange(ctx, pdClient, txnLeftBound, txnRightBound, client)
			if err != nil {
				log.Error("UnsafeDestroyRange txn mode range failed", zap.Error(err))
				return successed, err
			}
			record.ArchivedKeyspace = true
			successed = true
			log.Info("UnsafeDestroyRange success.", zap.Uint32("keyspaceId", keyspaceId))
		} else {
			log.Info("confirm yes but '-isrun' is not set, skip the operator.")
		}
	} else {
		log.Info("confirm: NO.")
	}
	return
}

func LoadKeyspaceAndDeleteRange(ctx context.Context, file *os.File,
	pdClient pd.Client, client *txnkv.Client, isRun, isSkipConfirm bool, keyspaceInterval time.Duration) error {
	recordsMap, err := common.CacheArchiveKeyspaceRecords(file)
	if err != nil {
		log.Error("Fialed to load keyspace records from file", zap.Error(err))
		return err
	}

	failedKsCt := 0
	successedKsCt := 0
	totalKsArchive := 0
	for _, record := range recordsMap {
		totalKsArchive++
		successed, err := DeleteRange(ctx, pdClient, client, record, isRun, isSkipConfirm)
		if err != nil {
			failedKsCt++
		}
		if successed {
			successedKsCt++
		}
		if totalKsArchive < len(recordsMap) {
			waitBetweenKeyspaces(keyspaceInterval)
		}
	}

	err = common.UpdateArchiveKeyspaceRecords(recordsMap, file)
	if err != nil {
		log.Error("Fialed to update KeyspaceRecordsFile", zap.Error(err))
		return err
	}
	log.Info("GC data ranges end.",
		zap.Int("totalKsArchive", totalKsArchive),
		zap.Int("failedKsCt", failedKsCt),
		zap.Int("successedKsCt", successedKsCt))
	return nil
}

// getStoresForGC gets the list of stores that needs to be processed during GC.
func getStoresForGC(ctx context.Context, pdClient pd.Client) ([]*metapb.Store, error) {
	stores, err := pdClient.GetAllStores(ctx)
	if err != nil {
		return nil, errors.Trace(err)
	}

	upStores := make([]*metapb.Store, 0, len(stores))
	for _, store := range stores {
		needsGCOp, err := needsGCOperationForStore(store)
		if err != nil {
			return nil, errors.Trace(err)
		}
		if needsGCOp {
			upStores = append(upStores, store)
		}
	}
	return upStores, nil
}

// needsGCOperationForStore checks if the store-level requests related to GC needs to be sent to the store.
// The store-level requests includes UnsafeDestroyRange, PhysicalScanLock, etc.
func needsGCOperationForStore(store *metapb.Store) (bool, error) {
	// TombStone means the store has been removed from the cluster and there isn't any peer on the store,
	// so needn't do GC for it. Offline means the store is being removed from the cluster and it becomes tombstone
	// after all peers are removed from it,
	// so we need to do GC for it.
	if store.State == metapb.StoreState_Tombstone {
		return false, nil
	}

	engineLabel := ""
	for _, label := range store.GetLabels() {
		if label.GetKey() == placement.EngineLabelKey {
			engineLabel = label.GetValue()
			break
		}
	}

	switch engineLabel {
	case placement.EngineLabelTiFlash, placement.EngineLabelTiFlashCompute:
		// For a TiFlash node, it uses other approach to delete dropped tables, so it's safe to skip sending
		// UnsafeDestroyRange requests; it has only learner peers and their data must exist in TiKV, so it's safe to
		// skip physical resolve locks for it.
		return false, nil

	case placement.EngineLabelTiKV, "":
		// If no engine label is set, it should be a TiKV node.
		return true, nil

	default:
		return true, errors.Errorf("unsupported store engine \"%v\" with storeID %v, addr %v",
			engineLabel,
			store.GetId(),
			store.GetAddress())
	}
}

// UnsafeDestroyRange Cleans up all keys in a range[startKey,endKey) and quickly free the disk space.
// The range might span over multiple regions, and the `ctx` doesn't indicate region. The request will be done directly
// on RocksDB, bypassing the Raft layer. User must promise that, after calling `UnsafeDestroyRange`,
// the range will never be accessed any more. However, `UnsafeDestroyRange` is allowed to be called
// multiple times on an single range.
func UnsafeDestroyRange(ctx context.Context, pdClient pd.Client,
	startKey []byte, endKey []byte, client *txnkv.Client) error {
	// Get all stores every time deleting a region. So the store list is less probably to be stale.
	stores, err := getStoresForGC(ctx, pdClient)
	if err != nil {
		metrics.TiKVUnsafeDestroyRangeFailuresCounterVec.WithLabelValues("get_stores").Inc()
		return err
	}

	req := tikvrpc.NewRequest(tikvrpc.CmdUnsafeDestroyRange, &kvrpcpb.UnsafeDestroyRangeRequest{
		StartKey: startKey,
		EndKey:   endKey,
	}, kvrpcpb.Context{DiskFullOpt: kvrpcpb.DiskFullOpt_AllowedOnAlmostFull})

	var wg sync.WaitGroup
	errChan := make(chan error, len(stores))

	for _, store := range stores {
		address := store.Address
		storeID := store.Id
		wg.Add(1)
		go func() {
			defer wg.Done()

			resp, err1 := client.GetTiKVClient().SendRequest(ctx, address, req, unsafeDestroyRangeTimeout)
			if err1 == nil {
				if resp == nil || resp.Resp == nil {
					err1 = errors.Errorf("[unsafe destroy range] returns nil response from store %v", storeID)
				} else {
					errStr := (resp.Resp.(*kvrpcpb.UnsafeDestroyRangeResponse)).Error
					if len(errStr) > 0 {
						err1 = errors.Errorf("[unsafe destroy range] range failed on store %v: %s", storeID, errStr)
					}
				}
			}

			if err1 != nil {
				metrics.TiKVUnsafeDestroyRangeFailuresCounterVec.WithLabelValues("send").Inc()
			}
			errChan <- err1
		}()
	}

	var errs []string
	for range stores {
		err1 := <-errChan
		if err1 != nil {
			errs = append(errs, err1.Error())
		}
	}

	wg.Wait()

	if len(errs) > 0 {
		return errors.Errorf("[unsafe destroy range] destroy range finished with errors: %v", errs)
	}
	return nil
}

func DumpKeyspaceList(ctx context.Context, metaServerAddr string, f *os.File,
	delayDuration time.Duration, expectedState keyspacepb.KeyspaceState, limit int) error {
	records, err := common.GetKeyspace(ctx, metaServerAddr, delayDuration, expectedState, limit)
	if err != nil {
		return err
	}
	return records.WriteTo(f)
}

func parseKeyspaceID(b []byte) uint32 {
	buf := append([]byte{}, b[:4]...)
	buf[0] = 0

	return binary.BigEndian.Uint32(buf)
}

func shrinkRegionIfNeeded(lower, upper []byte) ([]byte, []byte) {
	// region not contains txn keyspace, no need to shrink
	if bytes.Compare(txnKeyspaceStartBoundary, upper) >= 0 {
		return lower, upper
	}
	// region in txn keyspace, no need to shrink
	if bytes.Compare(lower, txnKeyspaceStartBoundary) >= 0 {
		return lower, upper
	}
	// region cross txn keyspace, shrink it
	return txnKeyspaceStartBoundary, upper
}

// CleanTxnRegionKeyspaceData cleans up all keys in a region in txn mode. If the region cross multiple keyspace,
// it will clean up all keys in keyspaces by issuing UnsafeDestroyRange requests keyspace by keyspace.
func CleanTxnRegionKeyspaceData(ctx context.Context, pdClient pd.Client, httpPdClient pdhttp.HttpPdClient,
	tikvClient *txnkv.Client, regionId uint64, isRun bool) error {
	// get region borders
	region, err := pdClient.GetRegionByID(ctx, regionId)
	if err != nil {
		log.Warn("Failed to get region by id", zap.Uint64("regionId", regionId), zap.Error(err))
	}
	if region == nil {
		log.Error("Region not exist", zap.Uint64("regionId", regionId))
		return nil
	}
	startKey, endKey := region.Meta.StartKey, region.Meta.EndKey

	// decode startKey & endKey
	_, rawStartKey, err := codec.DecodeBytes(startKey, nil)
	if err != nil {
		log.Error("Failed to decode startKey", zap.Error(err))
		return err
	}
	_, rawEndKey, err := codec.DecodeBytes(endKey, nil)
	if err != nil {
		log.Error("Failed to decode endKey", zap.Error(err))
		return err
	}
	// shrink region if needed
	rawStartKey, rawEndKey = shrinkRegionIfNeeded(rawStartKey, rawEndKey)

	if rawStartKey[0] != 'x' || rawEndKey[0] != 'x' {
		log.Error("Invalid region boundary",
			zap.Uint64("regionId", regionId),
			zap.String("startKey", hex.EncodeToString(rawStartKey)),
			zap.String("endKey", hex.EncodeToString(rawEndKey)))
		return nil
	}
	// calculate the keyspaces this region covers
	startKeyspace := parseKeyspaceID(rawStartKey)
	endKeyspace := parseKeyspaceID(rawEndKey)
	if len(rawEndKey) == keyspacePrefixLen {
		endKeyspace -= 1
	}

	log.Info("Start to clean region data",
		zap.Uint64("regionId", regionId),
		zap.Uint32("totalKeyspace", endKeyspace-startKeyspace+1),
		zap.String("startKey",
			hex.EncodeToString(rawStartKey)),
		zap.String("endKey", hex.EncodeToString(rawEndKey)))
	for id := startKeyspace; id <= endKeyspace; id++ {
		_, _, txnLeftBound, txnRightBound := common.GetRange(id)

		// confirm the keyspace has beed archived or deleted or not exists
		keyspaceNotExists := false
		keyspaceMeta, err := httpPdClient.GetKeyspaceByID(ctx, id)
		if err != nil {
			if pdhttp.IsKeyspaceNotExists(err) {
				keyspaceNotExists = true
			} else {
				log.Error("Failed to get keyspace meta, skip", zap.Uint32("keyspaceId", id), zap.Error(err))
				continue
			}
		}
		if !keyspaceNotExists && keyspaceMeta.State != "ARCHIVED" && keyspaceMeta.State != "TOMBSTONE" {
			log.Error("Keyspace state is not ARCHIVED or TOMBSTONE, skip",
				zap.Uint32("keyspaceId", id),
				zap.String("state", keyspaceMeta.State))
			continue
		}

		log.Info("UnsafeDestroyRange",
			zap.Uint32("keyspaceId", id),
			zap.String("txnLeftBound", hex.EncodeToString(txnLeftBound)),
			zap.String("txnRightBound", hex.EncodeToString(txnRightBound)))
		if !isRun {
			continue
		}
		if err = UnsafeDestroyRange(ctx, pdClient, txnLeftBound, txnRightBound, tikvClient); err != nil {
			log.Error("Failed to clean region data", zap.Uint64("regionId", regionId), zap.Error(err))
			return err
		}
	}
	return nil
}