// Copyright 2016-present, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package mocktikv_test

import (
	"bytes"
	"context"
	"math"
	"strconv"
	"time"

	. "github.com/pingcap/check"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/sessionctx/stmtctx"
	"github.com/pingcap/tidb/store/mockstore"
	"github.com/pingcap/tidb/store/mockstore/mocktikv"
	"github.com/pingcap/tidb/tablecodec"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/codec"
	"github.com/pingcap/tidb/util/rowcodec"
)

var _ = Suite(&testClusterSuite{})

type testClusterSuite struct {
	store kv.Storage
}

func (s *testClusterSuite) TestClusterSplit(c *C) {
	cluster := mocktikv.NewCluster()
	mocktikv.BootstrapWithSingleStore(cluster)
	mvccStore := mocktikv.MustNewMVCCStore()
	store, err := mockstore.NewMockTikvStore(
		mockstore.WithCluster(cluster),
		mockstore.WithMVCCStore(mvccStore),
	)
	c.Assert(err, IsNil)

	txn, err := store.Begin()
	c.Assert(err, IsNil)

	// Mock inserting many rows in a table.
	tblID := int64(1)
	idxID := int64(2)
	colID := int64(3)
	handle := int64(1)
	sc := &stmtctx.StatementContext{TimeZone: time.UTC}
	for i := 0; i < 1000; i++ {
		rowKey := tablecodec.EncodeRowKeyWithHandle(tblID, handle)
		colValue := types.NewStringDatum(strconv.Itoa(int(handle)))
		// TODO: Should use session's TimeZone instead of UTC.
		rd := rowcodec.Encoder{Enable: true}
		rowValue, err1 := tablecodec.EncodeRow(sc, []types.Datum{colValue}, []int64{colID}, nil, nil, &rd)
		c.Assert(err1, IsNil)
		txn.Set(rowKey, rowValue)

		encodedIndexValue, err1 := codec.EncodeKey(sc, nil, []types.Datum{colValue, types.NewIntDatum(handle)}...)
		c.Assert(err1, IsNil)
		idxKey := tablecodec.EncodeIndexSeekKey(tblID, idxID, encodedIndexValue)
		txn.Set(idxKey, []byte{'0'})
		handle++
	}
	err = txn.Commit(context.Background())
	c.Assert(err, IsNil)

	// Split Table into 10 regions.
	cluster.SplitTable(mvccStore, tblID, 10)

	// 10 table regions and first region and last region.
	regions := cluster.GetAllRegions()
	c.Assert(regions, HasLen, 12)

	allKeysMap := make(map[string]bool)
	recordPrefix := tablecodec.GenTableRecordPrefix(tblID)
	for _, region := range regions {
		startKey := mocktikv.MvccKey(region.Meta.StartKey).Raw()
		endKey := mocktikv.MvccKey(region.Meta.EndKey).Raw()
		if !bytes.HasPrefix(startKey, recordPrefix) {
			continue
		}
		pairs := mvccStore.Scan(startKey, endKey, math.MaxInt64, math.MaxUint64, kvrpcpb.IsolationLevel_SI, nil)
		if len(pairs) > 0 {
			c.Assert(pairs, HasLen, 100)
		}
		for _, pair := range pairs {
			allKeysMap[string(pair.Key)] = true
		}
	}
	c.Assert(allKeysMap, HasLen, 1000)

	cluster.SplitIndex(mvccStore, tblID, idxID, 10)

	allIndexMap := make(map[string]bool)
	indexPrefix := tablecodec.EncodeTableIndexPrefix(tblID, idxID)
	regions = cluster.GetAllRegions()
	for _, region := range regions {
		startKey := mocktikv.MvccKey(region.Meta.StartKey).Raw()
		endKey := mocktikv.MvccKey(region.Meta.EndKey).Raw()
		if !bytes.HasPrefix(startKey, indexPrefix) {
			continue
		}
		pairs := mvccStore.Scan(startKey, endKey, math.MaxInt64, math.MaxUint64, kvrpcpb.IsolationLevel_SI, nil)
		if len(pairs) > 0 {
			c.Assert(pairs, HasLen, 100)
		}
		for _, pair := range pairs {
			allIndexMap[string(pair.Key)] = true
		}
	}
	c.Assert(allIndexMap, HasLen, 1000)
}