package main

import (
	"database/sql"
	"fmt"
	"math/rand"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger-test/dmctl/pb"
	data "github.com/pingcap/schrodinger-test/pkg/ddl_data"
	"github.com/pingcap/schrodinger-test/util"
	"github.com/pingcap/schrodinger/pkg/test_client"
	"golang.org/x/net/context"
)

// ControllOptions is used to pass options from outside
type ControllOptions struct {
	mode              string
	failResume        bool
	masterSlaveSwitch bool
}

type DMController struct {
	ctx     context.Context
	cancel  context.CancelFunc
	nemesis *nemesis
	ddls    map[int][]string // caseID -> ddl data list
	errCh   chan error
	tidb    *sql.DB
	opts    *ControllOptions

	// whether DM-worker is alive and we can call `query-status` or `query-error` from nemesis.dmctl
	queryEnable struct {
		sync.Mutex
		enabled bool
	}
}

// CheckFullStageDone checks whether all subtasks' full stage is done
func (dmc *DMController) CheckFullStageDone(workerCount, caseCount int, fullDone *int32) {
	if dmc.opts.mode != AllMode {
		return
	}
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-dmc.ctx.Done():
			return
		case <-ticker.C:
			units, err := dmc.nemesis.dmctl.queryTasksUnit()
			if err != nil {
				dmc.errCh <- errors.Trace(err)
				return
			}
			ready := true
			count := 0
			for worker, m := range units {
				client.Debugf("worker: %s units stage: %+v", worker, m)
				for _, unit := range m {
					if unit == pb.UnitType_Check || unit == pb.UnitType_Load || unit == pb.UnitType_Dump {
						ready = false
					}
					count++
				}
			}
			if ready && count == workerCount*caseCount {
				atomic.CompareAndSwapInt32(fullDone, 0, 1)
				return
			}
		}
	}
}

// Run is DMController's main loop
func (dmc *DMController) Run() {
	if dmc.opts.failResume {
		err := dmc.addSkipSQLPatterns(false)
		if err != nil {
			dmc.nemesis.errCh <- errors.Trace(err)
			return
		}
		go dmc.autoReplaceSQLs()
	}
	for {
		var err error
		select {
		case err = <-dmc.errCh:
			dmc.nemesis.errCh <- errors.Trace(err)
			return
		case <-dmc.ctx.Done():
			return
		case <-time.After(time.Second * 5):
		}

		if dmc.opts.mode == AllMode {
			err = dmc.allModeRound()
		} else if dmc.opts.mode == FullMode {
			err = dmc.fullModeRound()
		} else {
			err = errors.Errorf("unknown task mode: %s", dmc.opts.mode)
		}
		if err != nil {
			dmc.nemesis.errCh <- errors.Trace(err)
			return
		}
	}
}

func (dmc *DMController) fullModeRound() error {
	switch rand.Intn(6) {
	case 1:
		dmc.nemesis.dmlMu.Lock()
		for caseID, st := range dmc.nemesis.dmlStat {
			if st == 0 {
				select {
				case <-dmc.nemesis.dmlCh[caseID]:
				default:
				}
				dmc.nemesis.dmlCh[caseID] <- true
				dmc.nemesis.dmlStat[caseID] = 1
			}
		}
		dmc.nemesis.dmlMu.Unlock()

	case 2:
		dmc.nemesis.checkMu.Lock()
		for caseID, st := range dmc.nemesis.checkStat {
			if st == 0 {
				select {
				case <-dmc.nemesis.checkCh[caseID]:
				default:
				}
				dmc.nemesis.checkCh[caseID] <- true
				dmc.nemesis.checkStat[caseID] = 1
			}
		}
		dmc.nemesis.checkMu.Unlock()
	default:
		dmc.nemesis.emptyMu.Lock()
		for _, ch := range dmc.nemesis.emptyCh {
			select {
			case <-ch:
			default:
			}
			ch <- true
		}
		dmc.nemesis.emptyMu.Unlock()
	}

	return nil
}

func (dmc *DMController) allModeRound() error {
	switch rand.Intn(6) {
	case 0, 4, 5:
		client.Debug("[control] empty round")
		dmc.nemesis.emptyMu.Lock()
		for _, ch := range dmc.nemesis.emptyCh {
			select {
			case <-ch:
			default:
			}
			ch <- true
		}
		dmc.nemesis.emptyMu.Unlock()

	case 1:
		client.Debug("[control] dml round")

		ready := dmc.checkKillWorkerReady()
		if ready {
			// safe wait for sharding DDL done, etc
			time.Sleep(time.Minute)
			dmc.updateQueryEnable(false)
			err := dmc.killGuys()
			if err != nil {
				return errors.Trace(err)
			}
			if dmc.opts.failResume {
				err = dmc.addSkipSQLPatterns(true)
				if err != nil {
					return errors.Trace(err)
				}
			}
			dmc.updateQueryEnable(true)
		}

		dmc.nemesis.dmlMu.Lock()
		for caseID, st := range dmc.nemesis.dmlStat {
			if st == 0 {
				select {
				case <-dmc.nemesis.dmlCh[caseID]:
				default:
				}
				dmc.nemesis.dmlCh[caseID] <- true
				dmc.nemesis.dmlStat[caseID] = 1
			}
		}
		dmc.nemesis.dmlMu.Unlock()

	case 2:
		client.Debug("[control] check round")
		dmc.nemesis.checkMu.Lock()
		for caseID, st := range dmc.nemesis.checkStat {
			if st == 0 {
				select {
				case <-dmc.nemesis.checkCh[caseID]:
				default:
				}
				dmc.nemesis.checkCh[caseID] <- true
				dmc.nemesis.checkStat[caseID] = 1
			}
		}
		dmc.nemesis.checkMu.Unlock()

	case 3:
		client.Debug("[control] master slave switch round")
		if !dmc.opts.masterSlaveSwitch {
			return nil
		}
		err := dmc.upstreamMasterSlaveSwitch()
		if err != nil {
			return errors.Trace(err)
		}
	}

	return nil
}

func NewDMController(ctx context.Context, nemesis *nemesis, opts *ControllOptions, ddls map[int][]string) *DMController {
	cctx, cancel := context.WithCancel(ctx)
	return &DMController{
		ctx:     cctx,
		cancel:  cancel,
		nemesis: nemesis,
		opts:    opts,
		ddls:    ddls,
		errCh:   make(chan error),
	}
}

func (dmc *DMController) Init(cli *client.Client) error {
	tidbIP := strings.Split(cli.TiDBService(), ":")[0]
	db, err := util.OpenDB(fmt.Sprintf("root:@tcp(%s:4000)/", tidbIP), 20)
	if err != nil {
		return errors.Trace(err)
	}
	dmc.tidb = db
	dmc.updateQueryEnable(true)
	return nil
}

// stop and restart a random dm-worker
func (dmc *DMController) killGuys() error {
	rid := rand.Intn(len(dmc.nemesis.workers))
	return errors.Trace(dmc.stopStartWorker(dmc.nemesis.workers[rid]))
}

func (dmc *DMController) restartWorker(index int) error {
	workerIP := dmc.nemesis.workers[index]
	client.Infof("stop dm worker..%s", workerIP)
	var err error
	err = dmc.nemesis.stopWorker(workerIP, DOWN)
	if err != nil {
		return errors.Trace(err)
	}

	client.Infof("start dm-worker...%s", workerIP)
	err = dmc.nemesis.startWorkerOnly(workerIP)
	if err != nil {
		return errors.Trace(err)
	}
	status, _ := dmc.nemesis.dmctl.queryStatus()
	client.Infof("query status %s", status)
	return errors.Trace(err)
}

func (dmc *DMController) stopStartWorker(workerIP string) error {
	client.Infof("stop dm worker..%s", workerIP)

	var err error
	switch rand.Intn(2) {
	case 0:
		err = dmc.nemesis.stopWorker(workerIP, KILL)
	case 1:
		err = dmc.nemesis.stopWorker(workerIP, DOWN)
	}
	if err != nil {
		return errors.Trace(err)
	}

	client.Infof("start dm-worker...%s", workerIP)
	err = dmc.nemesis.startWorkerOnly(workerIP)
	if err != nil {
		return errors.Trace(err)
	}

	status, _ := dmc.nemesis.dmctl.queryStatus()
	client.Infof("task status %s", status)

	return nil
}

func (dmc *DMController) checkKillWorkerReady() bool {
	dmc.nemesis.queryG.request.Lock()
	requestIDs := make([]int, 0, len(dmc.nemesis.queryG.request.ch))
	for caseID, ch := range dmc.nemesis.queryG.request.ch {
		select {
		case <-ch:
		default:
		}
		ch <- &Message{t: msgSyncCheckpoint}
		requestIDs = append(requestIDs, caseID)
	}
	dmc.nemesis.queryG.request.Unlock()

	var wg sync.WaitGroup
	chs := make(map[int]chan *MessageResponse)
	dmc.nemesis.queryG.response.Lock()
	for _, caseID := range requestIDs {
		chs[caseID] = dmc.nemesis.queryG.response.ch[caseID]
	}
	dmc.nemesis.queryG.response.Unlock()

	response := make(chan *MessageResponse, len(chs))
	for caseID, ch := range chs {
		wg.Add(1)
		go func(id int, rch chan *MessageResponse) {
			defer wg.Done()
			select {
			case resp := <-rch:
				response <- resp
			case <-time.After(time.Minute * 5):
				client.Warnf("wait for case-%d responseCh timeout", id)
			}
		}(caseID, ch)
	}
	wg.Wait()

	finished := true
	for len(response) > 0 {
		resp := <-response
		for dbIndex, ready := range resp.syncCheckpoint {
			if !ready {
				client.Infof("case-%d db-%d syncer checkpoint not saved", resp.caseID, dbIndex)
				finished = false
			}
		}
	}
	return finished
}

func (dmc *DMController) upstreamMasterSlaveSwitch() error {
	var wg sync.WaitGroup

	dbIndex := rand.Intn(len(dmc.nemesis.workers))

	// step1: send prepare to all dm cases
	prevote := 0
	dmc.nemesis.mssG.precheck.Lock()
	for _, ch := range dmc.nemesis.mssG.precheck.ch {
		select {
		case <-ch:
		default:
		}
		ch <- dbIndex
		prevote++
	}
	dmc.nemesis.mssG.precheck.Unlock()

	if prevote == 0 {
		return nil
	}
	client.Debugf("send %d prevote to start master slave switch", prevote)

	// step2: wait every dm case ready
	chs := make(map[int]chan bool)
	dmc.nemesis.mssG.ready.Lock()
	for caseID, ch := range dmc.nemesis.mssG.ready.ch {
		chs[caseID] = ch
	}
	dmc.nemesis.mssG.ready.Unlock()

	var ready int32
	var notReady int32
	for caseID, ch := range chs {
		wg.Add(1)
		go func(id int, rch chan bool) {
			defer wg.Done()
			select {
			case isReady := <-rch:
				if isReady {
					atomic.AddInt32(&ready, 1)
				} else {
					atomic.AddInt32(&notReady, 1)
				}
			case <-time.After(time.Minute * 5):
				atomic.AddInt32(&notReady, 1)
				client.Warnf("wait for case-%d readyCh timeout", id)
			}
		}(caseID, ch)
	}
	wg.Wait()

	if int(ready+notReady) != prevote {
		return errors.Errorf("[dm] prevote %d not equals to ready %d during masterSlaveSwitch", prevote, (ready + notReady))
	}
	client.Debugf("receive %d ready for master slave switch", ready)

	// step3: tell each case to do master slave switch staff, such as update
	// db connections in executor, leader does some real switching staff.
	dmc.nemesis.mssG.switch1.Lock()
	leader := true
	for _, ch := range dmc.nemesis.mssG.switch1.ch {
		ch <- &mssOptions{
			leader:  leader,
			dbIndex: dbIndex,
			abort:   notReady > 0,
		}
		if leader {
			leader = false
		}
	}
	dmc.nemesis.mssG.switch1.Unlock()

	if notReady > 0 {
		client.Infof("%d cases not ready, abort master slave switch", notReady)
		return nil
	}

	client.Debugf("send %d switch1 for master slave switch", len(dmc.nemesis.mssG.switch1.ch))

	// step4: wait all dm cases master slave switch staff completed
	chs = make(map[int]chan bool)
	dmc.nemesis.mssG.switch2.Lock()
	for caseID, ch := range dmc.nemesis.mssG.switch2.ch {
		chs[caseID] = ch
	}
	dmc.nemesis.mssG.switch2.Unlock()

	var switch2 int32
	for caseID, ch := range chs {
		wg.Add(1)
		go func(id int, rch chan bool) {
			defer wg.Done()
			select {
			case <-rch:
				atomic.AddInt32(&switch2, 1)
			case <-time.After(time.Minute * 10):
				dmc.errCh <- errors.Errorf("wait for case-%d switch2 timeout, received %d", caseID, switch2)
			}
		}(caseID, ch)
	}
	wg.Wait()

	if ready != switch2 {
		return errors.Errorf("[dm] switch %d not equals to ready %d during masterSlaveSwitch", switch2, ready)
	}

	// restart worker
	err := dmc.restartWorker(dbIndex)
	if err != nil {
		return errors.Trace(err)
	}

	// step5: notify all test cases switch done
	dmc.nemesis.mssG.done.Lock()
	for _, ch := range dmc.nemesis.mssG.done.ch {
		select {
		case <-ch:
		default:
		}
		ch <- true
	}
	dmc.nemesis.mssG.done.Unlock()

	return nil
}

// we have two kinds of incompatible SQLs in test case:
//   1. chagne varchar to a shorter varchar
//   2. drop column with index
// as we can skip/replace incompatible SQLs in either sharding or non-sharding
// mode, and operator can be applied before or after incompatible SQLs executed,
// we have 8 different test combinations. We only test the following four scenes:
//   1. change to shorter varchar,	sharding mode,		skip,		before executed,	use sql pattern
//   2. change to shorter varchar,	non-sharding mode,	skip,		before executed,	use sql pattern
//   3. drop column with index,		sharding mode,		replace,	after executed,		use binlog position
//   4. drop column with index,		non-sharding mode,	replace,	after executed,		use binlog position
// We use addSkipSQLPatterns and autoReplaceSQLs to achieve these four test cases

// addSkipSQLPatterns pre add sql replace operator for addSkipSQLPatterns SQLs
// onRestart represents whether this function is called after a DM-worker restart,
// if on Restart is true, we only send sql_skip_operation in non-sharding test cases
// otherwise we send all sql_skip_operations
func (dmc *DMController) addSkipSQLPatterns(onRestart bool) error {
	for caseID, ddls := range dmc.ddls {
		if onRestart && !isNonShardingCase(caseID) {
			continue
		}
		for _, ddl := range ddls {
			details := strings.Split(ddl, " ")
			ddlType, err := strconv.Atoi(details[0])
			if err != nil {
				return errors.Annotatef(err, "invalid ddl type: %s", ddlType)
			}
			// currently in DM test cases we don't enable ALTER TABLE COLUMN DDL
			// by default, this ALTER TABLE COLUMN DDL comes from the custom
			// AddAndAlterColumn DDL type.
			if ddlType == data.AlterColumn && details[2] == "VARCHAR(20)" {
				var (
					worker      string
					sharding    bool
					taskName    = genTaskName(caseID)
					basePattern = "~(?i)ALTER\\s+TABLE\\s+`.*`.`%s`\\s+MODIFY\\s+COLUMN\\s+`%s`\\s+VARCHAR\\(20\\)"
				)
				if isNonShardingCase(caseID) {
					// non sharding mode must skip in each table
					for idx, w := range dmc.nemesis.workers {
						var tbls []string
						worker = fmt.Sprintf("%s:%d", w, DMWorkerPort)
						if idx == 0 {
							tbls = totalSkipTables1
						} else {
							// except for the first worker, other subtask only has one table router tf_x -> tf_x
							tbls = []string{fmt.Sprintf("tf_%d", (idx + 1))}
						}
						for _, tbl := range tbls {
							sqlPattern := fmt.Sprintf(basePattern, tbl, details[1])
							err := dmc.nemesis.dmctl.sqlSkipPattern(taskName, worker, sqlPattern, sharding)
							if err != nil {
								return errors.Trace(err)
							}
						}
					}
				} else {
					sqlPattern := fmt.Sprintf(basePattern, ".*", details[1])
					sharding = true
					err := dmc.nemesis.dmctl.sqlSkipPattern(taskName, worker, sqlPattern, sharding)
					if err != nil {
						return errors.Trace(err)
					}
				}
			}
		}
	}
	return nil
}

// autoResumeFromWarning runs in background and auto replaces incompatible SQLs
func (dmc *DMController) autoReplaceSQLs() {
	for {
		select {
		case <-dmc.ctx.Done():
			return
		case <-time.After(1 * time.Second):
			dmc.queryEnable.Lock()
			if dmc.queryEnable.enabled {
				err := dmc.ReplaceSQLAndResumeTask()
				if err != nil {
					dmc.errCh <- err
					dmc.queryEnable.Unlock()
					return
				}
			}
			dmc.queryEnable.Unlock()
		}
	}
}

func (dmc *DMController) genReplaceSQLs(errSQL string) ([]string, error) {
	result := make([]string, 0)
	r := regexp.MustCompile(".*ALTER TABLE `(?P<schema>.+)`.`(?P<table>.+)` DROP COLUMN `(?P<column>.+)`.*")
	found := r.FindStringSubmatch(errSQL)
	if len(found) != 4 {
		return nil, errors.Errorf("invalid error SQL: %s", errSQL)
	}
	var (
		schema = found[1]
		table  = found[2]
		column = found[3]
	)
	indexName, err := getColumnIndexName(dmc.tidb, schema, table, column)
	if err != nil {
		return nil, errors.Trace(err)
	}
	result = append(result, fmt.Sprintf("ALTER TABLE `%s`.`%s` DROP INDEX %s", schema, table, indexName))
	result = append(result, fmt.Sprintf("ALTER TABLE `%s`.`%s` DROP COLUMN `%s`", schema, table, column))

	return result, nil
}

// ReplaceSQLAndResumeTask checks task warning, skips abnormal sql and resumes paused task
func (dmc *DMController) ReplaceSQLAndResumeTask() error {
	syncErrors, err := dmc.nemesis.dmctl.querySyncError()
	if err != nil {
		return errors.Trace(err)
	}
	// to be resumed task name -> belonging worker
	resumes := make(map[string][]string)
	for worker, posList := range syncErrors {
		for taskName, serr := range posList {
			replaceSQLs, err := dmc.genReplaceSQLs(serr.errorSQL)
			if err != nil {
				return errors.Trace(err)
			}
			if err := dmc.nemesis.dmctl.sqlReplaceWithPos(taskName, worker, serr.binlogPos, []string{strings.Join(replaceSQLs, ";")}); err != nil {
				return errors.Trace(err)
			}
			if _, ok := resumes[taskName]; !ok {
				resumes[taskName] = make([]string, 0)
			}
			resumes[taskName] = append(resumes[taskName], worker)
		}
	}
	for taskName, workers := range resumes {
		if err := dmc.nemesis.dmctl.resumeTask(taskName, workers); err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func (dmc *DMController) updateQueryEnable(enable bool) {
	dmc.queryEnable.Lock()
	defer dmc.queryEnable.Unlock()
	dmc.queryEnable.enabled = enable
}