// Copyright 2017 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 main

import (
	"context"
	"database/sql"
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger-test/util"
	"github.com/pingcap/schrodinger/pkg/test_client"
)

var (
	createSQL = `DROP TABLE IF EXISTS merge_test%d;
	CREATE TABLE merge_test%d (
		id bigint,
		flag bigint,
		timestamp bigint,
		reserved1 text,
		reserved2 text,
		reserved3 text,
		reserved4 text,
		reserved5 text,
		reserved6 text,
		reserved7 text,
		reserved8 text,
		reserved9 text,
		reserved10 text,
		reserved11 text,
		reserved12 text,
		reserved13 text,
		reserved14 text,
		reserved15 text,
		reserved16 text,
		reserved17 text,
		reserved18 text,
		reserved19 text,
		reserved20 text,
		reserved21 text,
		reserved22 text,
		reserved23 text,
		reserved24 text,
		reserved25 text,
		reserved26 text,
		reserved27 text,
		reserved28 text,
		reserved29 text,
		reserved30 text
	);`
)

const retryLimit = 5
const concurrency = 500

var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func randSeq(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

// DB is db
type DB struct {
	*sql.DB
	id   int
	flag int
}

// InitDB inits db
func InitDB(dbAddr, dbName string) (*DB, error) {
	var err error
	dbDSN := fmt.Sprintf("root:@tcp(%s)/%s", dbAddr, dbName)
	db, err := util.OpenDB(dbDSN, 20)
	if err != nil {
		return nil, errors.Trace(err)
	}

	return &DB{DB: db, id: 1, flag: 1}, nil
}

// InsertUntilSplitToRegions inserts until split to regines
func (db *DB) InsertUntilSplitToRegions(tableID int64, n int) error {
	ctx, cancel := context.WithCancel(context.Background())
	wg := db.Insert(ctx, db.id-1, 100)
	defer wg.Wait()
	defer cancel()

	for {
		regions, err := GetRegions()
		if err != nil {
			return errors.Trace(err)
		}

		find := 0
		for _, region := range regions {
			if GetTableID(region) == tableID {
				find++
			}
		}

		client.Infof("find %d regions for table %d", find, tableID)
		if find >= n {
			break
		}

		time.Sleep(time.Second * 60)
	}
	client.Info("------------------------------------------------")

	return nil
}

// WaitUntilMerge print table regions info periodically until all merged.
func (db *DB) WaitUntilMerge(tableRegions []*RegionInfo) error {
	tableID := GetTableID(tableRegions[0])
	for _, region := range tableRegions {
		if GetTableID(region) != tableID {
			return fmt.Errorf("unexpected table regions")
		}
	}

	for {
		time.Sleep(time.Second * 60)

		isMerge := false
		for i := len(tableRegions) - 1; i >= 0; i-- {
			region := tableRegions[i]
			nowRegion, err := GetRegionInfo(region.ID)
			if err != nil {
				return errors.Trace(err)
			}

			if nowRegion.ID == 0 {
				isMerge = true
				tableRegions = append(tableRegions[:i], tableRegions[i+1:]...)
				client.Infof("region %d was merged", region.ID)
				continue
			}

			if GetTableID(nowRegion) != tableID {
				return fmt.Errorf("region %d's tableID has changed", region.ID)
			}

			client.Infof("[region %d] approximate size is %d", nowRegion.ID, nowRegion.ApproximateSize)
			tableRegions[i] = nowRegion
		}
		client.Info("------------------------------------------------")
		if isMerge && len(tableRegions) <= 1 {
			return nil
		}

		var err error
		tableRegions, err = CheckRegionsTableID(tableID, tableRegions)
		if err != nil {
			return errors.Trace(err)
		}
	}
}

// WaitUntilTimeout print table regions info periodically until all merged or timeout.
func (db *DB) WaitUntilTimeout(tableRegions []*RegionInfo, timeout time.Duration) error {
	tableID := GetTableID(tableRegions[0])
	for _, region := range tableRegions {
		if GetTableID(region) != tableID {
			return fmt.Errorf("unexpected table regions")
		}
	}

	for _, region := range tableRegions {
		nowRegion, err := GetRegionInfo(region.ID)
		if err != nil {
			return errors.Trace(err)
		}
		client.Infof("[region %d] approximate size is %d", nowRegion.ID, nowRegion.ApproximateSize)
	}
	client.Info("------------------------------------------------")

	time.Sleep(timeout)

	for _, region := range tableRegions {
		nowRegion, err := GetRegionInfo(region.ID)
		if err != nil {
			return errors.Trace(err)
		}

		if nowRegion.ID == 0 {
			client.Infof("[region %d] was merged", region.ID)
			continue
		}
		client.Infof("[region %d] approximate size is %d", nowRegion.ID, nowRegion.ApproximateSize)
	}
	client.Info("------------------------------------------------")

	return nil
}

// InsertIntoNewTable InsertIntoNewTable
func (db *DB) InsertIntoNewTable(ctx context.Context, batchSize int) {
	wg := db.Insert(ctx, db.id-1, batchSize)
	wg.Wait()
}

// Insert insert batch row into table.
func (db *DB) Insert(ctx context.Context, tableID int, batchSize int) *sync.WaitGroup {
	var wg sync.WaitGroup
	db.flag = 1
	for i := 0; i < concurrency; i++ {
		go func(ctx context.Context) {
			wg.Add(1)
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				default:
				}

				err := db.insertBatch(tableID, batchSize)
				if err != nil {
					client.Fatal(err)
				}
				db.flag++
			}
		}(ctx)
	}

	return &wg
}

func (db *DB) insertBatch(tableID int, batchSize int) error {
	now := time.Now().Unix()
	table := fmt.Sprintf("merge_test%d", tableID)

	values := make([]string, batchSize)
	for i := 0; i < batchSize; i++ {
		values[i] = fmt.Sprintf(`(%d,%d,%d`, rand.Int31n(100000), db.flag, now)
		for j := 1; j <= 30; j++ {
			values[i] += fmt.Sprintf(",\"%s\"", randSeq(255))
		}
		values[i] += ")"
	}

	sql := fmt.Sprintf(`INSERT INTO %s VALUES %s;`, table, strings.Join(values, ","))
	err := util.RunWithRetry(context.Background(), retryLimit, 3*time.Second, func() error {
		_, err := db.Exec(sql)
		return err
	})
	return errors.Trace(err)
}

func (db *DB) deleteBatch() error {
	sql := fmt.Sprintf(`DELETE FROM merge_test%d WHERE flag=%d;`, db.id-1, db.flag)
	err := util.RunWithRetry(context.Background(), retryLimit, 3*time.Second, func() error {
		_, err := db.Exec(sql)
		return err
	})
	return errors.Trace(err)
}

// DeleteAll delete all rows in table.
func (db *DB) DeleteAll() error {
	err := db.AddDeleteIndex()
	if err != nil {
		return errors.Trace(err)
	}
	for db.flag > 1 {
		db.flag--
		if err := db.deleteBatch(); err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

// DeletePartial delete half of rows in table.
func (db *DB) DeletePartial() error {
	err := db.AddDeleteIndex()
	if err != nil {
		return errors.Trace(err)
	}
	for db.flag > 1 {
		db.flag -= 2
		if err := db.deleteBatch(); err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

// DropTable drops table.
func (db *DB) DropTable() error {
	sql := fmt.Sprintf(`DROP TABLE merge_test%d;`, db.id-1)
	err := util.RunWithRetry(context.Background(), retryLimit, 3*time.Second, func() error {
		_, err := db.Exec(sql)
		return err
	})
	return errors.Trace(err)
}

// TruncateTable truncate table.
func (db *DB) TruncateTable() error {
	sql := fmt.Sprintf(`TRUNCATE TABLE merge_test%d`, db.id-1)
	err := util.RunWithRetry(context.Background(), retryLimit, 3*time.Second, func() error {
		_, err := db.Exec(sql)
		return err
	})

	return errors.Trace(err)
}

// CreateTable create a new table and return table id.
func (db *DB) CreateTable() (int64, error) {
	query := fmt.Sprintf(createSQL, db.id, db.id)
	db.id++
	err := util.RunWithRetry(context.Background(), retryLimit, 3*time.Second, func() error {
		_, err := db.Exec(query)
		return err
	})
	if err != nil {
		return 0, errors.Trace(err)
	}
	return GetTableIDByName(fmt.Sprintf("merge_test%d", db.id-1))
}

// AddIndex add index for table.
func (db *DB) AddIndex() error {
	sql := fmt.Sprintf(`ALTER TABLE merge_test%d ADD INDEX idx (reserved1(255));`, db.id-1)
	err := util.RunWithRetry(context.Background(), retryLimit, 3*time.Second, func() error {
		_, err := db.Exec(sql)
		return err
	})

	return errors.Trace(err)
}

// AddDeleteIndex add index for accerating delete.
func (db *DB) AddDeleteIndex() error {
	sql := fmt.Sprintf(`ALTER TABLE merge_test%d ADD INDEX deleteIdx (flag);`, db.id-1)
	_, err := db.Exec(sql)

	return errors.Trace(err)
}

// DropIndex drop index for table.
func (db *DB) DropIndex() error {
	sql := fmt.Sprintf(`DROP INDEX idx on merge_test%d;`, db.id-1)
	err := util.RunWithRetry(context.Background(), retryLimit, 3*time.Second, func() error {
		_, err := db.Exec(sql)
		return err
	})

	return errors.Trace(err)
}