// Copyright 2022 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 performance

import (
	"fmt"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/pingcap/endless/pkg/matcher"
	. "github.com/pingcap/endless/pkg/workload"
	. "github.com/pingcap/endless/testcase/geopartition/utils"
	"github.com/pingcap/log"
	"go.uber.org/zap"
)

const PartitionDateLayout = "2006-01-02"

type dtPartition struct {
	time.Time
	tableName string
}

func newPartition(table, dt string) dtPartition {
	return dtPartition{Time: Try(time.Parse(PartitionDateLayout, dt)).(time.Time), tableName: table}
}

func (p dtPartition) NextPartition() dtPartition {
	return dtPartition{Time: p.AddDate(0, 0, 1), tableName: p.tableName}
}

func (p dtPartition) Date() string {
	return p.Format(PartitionDateLayout)
}

func (p dtPartition) PartitionName() string {
	return strings.ReplaceAll("p"+p.Date(), "-", "")
}

func (p dtPartition) AddPartitionSQL() string {
	return fmt.Sprintf("ALTER TABLE %s ADD PARTITION (PARTITION %s VALUES IN ('%s'))",
		p.tableName, p.PartitionName(), p.Date())
}

func (p dtPartition) DropPartitionSQL() string {
	return fmt.Sprintf("ALTER TABLE %s DROP PARTITION %s", p.tableName, p.PartitionName())
}

func (p dtPartition) AlterToColdSQL() string {
	return fmt.Sprintf("ALTER TABLE %s PARTITION %s PLACEMENT POLICY cold",
		p.tableName, p.PartitionName())
}

func (p dtPartition) InsertDataSQL(afterIdx, cnt int) string {
	values := make([]string, 0, cnt)
	suffix := strings.Repeat("x", 1024)
	for i := 1; i <= cnt; i++ {
		text := fmt.Sprintf("text-%s-%d-%s", p.Date(), afterIdx+i, suffix)
		values = append(values, fmt.Sprintf("('%s', %d, '%s')", p.Date(), afterIdx+i, text))
	}
	return fmt.Sprintf("INSERT INTO %s VALUES %s", p.tableName, strings.Join(values, ","))
}

func frequencyPartitionDDLTest(ctx *SuiteContext) {
	It("Test frequency add/drop/alter partition for placement partition #frequency_partition_ddl_with_placement#", func() {
		tableName := "tp_frequency_perf"
		defer func() {
			MustExec(ctx.DB, "DROP TABLE IF EXISTS "+tableName)
			MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS ssd")
			MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS cold")
		}()
		MustExec(ctx.DB, "DROP TABLE IF EXISTS "+tableName)
		MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS ssd")
		MustExec(ctx.DB, "DROP PLACEMENT POLICY IF EXISTS cold")
		MustExec(ctx.DB, "CREATE PLACEMENT POLICY ssd CONSTRAINTS='[+disk=ssd]'")
		MustExec(ctx.DB, "CREATE PLACEMENT POLICY cold CONSTRAINTS='[-disk=ssd]'")

		MustExec(ctx.DB, "set @@session.tidb_enable_list_partition = ON")
		MustExec(ctx.DB, "CREATE TABLE "+tableName+" (d date NOT NULL, v int, c text) PLACEMENT POLICY ssd "+
			"PARTITION BY LIST COLUMNS(d) ("+
			"PARTITION p20000101 VALUES IN('2000-01-01')"+
			")")

		coldPartitions := make([]dtPartition, 0)
		hotPartitions := []dtPartition{newPartition(tableName, "2000-01-01")}
		maxHotPartitions := 90
		maxColdPartitions := 1000
		for {
			newPartition := hotPartitions[len(hotPartitions)-1].NextPartition()
			addSQL := newPartition.AddPartitionSQL()
			log.Info("Create new partition", zap.String("dt", newPartition.Date()), zap.String("SQL", addSQL))
			MustExec(ctx.DB, addSQL)
			hotPartitions = append(hotPartitions, newPartition)

			var move []dtPartition
			if len(hotPartitions) > maxHotPartitions {
				move = hotPartitions[0 : len(hotPartitions)-maxHotPartitions]
				for _, p := range move {
					alterSQL := p.AlterToColdSQL()
					log.Info("Alter partition to cold", zap.String("dt", p.Date()), zap.String("SQL", alterSQL))
					MustExec(ctx.DB, alterSQL)
				}
				coldPartitions = append(coldPartitions, move...)
				hotPartitions = hotPartitions[len(move):]
			}

			for i := 0; i < 10; i++ {
				MustExec(ctx.DB, newPartition.InsertDataSQL(i*1000, 1000))
			}

			if len(coldPartitions) > maxColdPartitions {
				drop := coldPartitions[0 : len(coldPartitions)-maxColdPartitions]
				for _, p := range drop {
					dropSQL := p.DropPartitionSQL()
					log.Info("Drop partition", zap.String("dt", p.Date()), zap.String("SQL", dropSQL))
					MustExec(ctx.DB, dropSQL)
				}
				coldPartitions = coldPartitions[len(drop):]
			}
		}
	})
}