package main

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

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

type controller struct {
	ctx context.Context
	sync.RWMutex
	nemesis *nemesis
	tidb    string

	errCh chan error
	done  chan struct{}
}

func newController(ctx context.Context, nemesisCli *nemesis, tidb string) *controller {
	return &controller{
		ctx:     ctx,
		nemesis: nemesisCli,
		tidb:    tidb,
		errCh:   make(chan error, 1),
		done:    make(chan struct{}, 1),
	}
}

func (c *controller) startNemesis() {
	client.Infof("run nemesis")

	ticker := time.NewTicker(10 * time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			err := c.runNemesisOnce()
			if err != nil {
				c.errCh <- err
				return
			}
		}
	}
}

func (c *controller) runNemesisOnce() error {
	// keep 1/3 times does nothing..
	c.Lock()
	defer c.Unlock()

	state, err := util.ProcessState(c.nemesis.lightning, "tidb-lightning")
	if err != nil {
		return errors.Trace(err)
	}

	if state == daemon.ProcStatNormalExited {
		err = c.checkCounts()
		if err == nil {
			c.done <- struct{}{}
			return nil
		}
		return errors.Trace(err)
	}

	switch rand.Intn(3) {
	case 0:
		return errors.Trace(c.stopStartImporter())
	case 1:
		return errors.Trace(c.stopStartLightning())
	}
	return nil
}

func (c *controller) check() {
	client.Infof("run check")
	ticker := time.NewTicker(20 * time.Minute)
	defer ticker.Stop()
	for {
		select {
		case <-c.ctx.Done():
			return
		case <-ticker.C:
			client.Infof("start check..")
			err := c.checkOnce()
			if err != nil {
				c.errCh <- err
				return
			}
		}
	}
}

func (c *controller) checkOnce() error {
	c.Lock()
	defer c.Unlock()

	state, err := util.ProcessState(c.nemesis.lightning, "tidb-lightning")
	if err != nil {
		return errors.Trace(err)
	}
	if state == daemon.ProcStatErrorExited {
		return errors.New("lightning exit with error")
	}
	if state == daemon.ProcStatNormalExited {
		err = c.checkCounts()
		if err == nil {
			c.done <- struct{}{}
			return nil
		}
		return errors.Trace(err)
	}

	state, err = util.ProcessState(c.nemesis.importer, "tikv-importer")
	if err != nil {
		return errors.Trace(err)
	}
	if state == daemon.ProcStatErrorExited {
		return errors.New("importer exit with error")
	}
	return nil
}

func (c *controller) checkCounts() error {
	db, err := util.OpenDB(fmt.Sprintf("root:@tcp(%s:4000)/schr", c.tidb), 20)
	if err != nil {
		return errors.Trace(err)
	}

	for i := 1; i <= 50; i++ {
		err := c.checkTableCount(db, fmt.Sprintf("sbtest%d", i), 10000000)
		if err != nil {
			return errors.Trace(err)
		}
	}

	return nil
}

func (c *controller) checkTableCount(db *sql.DB, tableName string, count int64) error {
	cnt, err := dbutil.GetRowCount(c.ctx, db, "schr", tableName, "")
	if err != nil {
		return errors.Trace(err)
	}

	if cnt != count {
		return errors.Errorf("count not match in %s, expect %d but got %d", tableName, count, cnt)
	}
	return nil
}

func (c *controller) stopStartImporter() error {
	client.Debugf("stop start importer..")

	time.Sleep(60 * time.Second)

	var mode stopMode
	switch rand.Intn(2) {
	case 0:
		mode = KILL
	case 1:
		mode = DOWN
	}

	client.Debugf("stop importer using mode %s", mode)

	err := c.nemesis.stopImporter(mode)
	if err != nil {
		return errors.Trace(err)
	}

	time.Sleep(2 * time.Minute)
	client.Debugf("start importer")
	return errors.Trace(c.nemesis.startImporter())
}

func (c *controller) stopStartLightning() error {
	client.Debugf("stop start lightning..")
	time.Sleep(60 * time.Second)

	var mode stopMode
	switch rand.Intn(2) {
	case 0:
		mode = KILL
	case 1:
		mode = DOWN
	}

	client.Debugf("stop lightning using mode %s", mode)

	err := c.nemesis.stopLightning(mode)

	if err != nil {
		return errors.Trace(err)
	}

	time.Sleep(2 * time.Minute)
	client.Debugf("start lightning")
	return errors.Trace(c.nemesis.startLightning())
}
