package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/pingcap/schrodinger-test/util"
	"github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/store/tikv"
	"github.com/pingcap/tidb/tablecodec"
	"github.com/pingcap/tidb/util/codec"
)

type ScanEndLimitCase struct {
	cli          *client.Client
	dbName       string
	timeLimit    time.Duration
	logInterval  time.Duration
	startTime    time.Time
	queryThreads int
}

func NewScanEndLimitCase(cli *client.Client, dbName string, timeLimit time.Duration, logInterval time.Duration, queryThreads int) *ScanEndLimitCase {
	return &ScanEndLimitCase{
		cli:          cli,
		dbName:       dbName,
		timeLimit:    timeLimit,
		logInterval:  logInterval,
		queryThreads: queryThreads,
	}
}

// Prepares for the test.
func (c *ScanEndLimitCase) Prepare() error {
	dbAddr := c.cli.TiDBService()
	dbDSN := fmt.Sprintf("root:@tcp(%s)/%s", dbAddr, *dbName)
	db, err := util.OpenDB(dbDSN, *concurrency)
	if err != nil {
		client.Errorf("[scan end limit] OpenDB error %v", errors.Trace(err))
		return errors.Trace(err)
	}
	defer db.Close()

	// Create a table `t1` with 100 rows.
	_, err = db.Exec("create table t1 (id integer, i integer, v varchar(256), primary key(id));")
	if err != nil {
		client.Errorf("[scan end limit] Failed preparing table: %v", errors.Trace(err))
		return errors.Trace(err)
	}
	for i := 0; i < 100; i++ {
		_, err = db.Exec("insert into t1 values (?, ?, ?);", i, i, randStr())
		if err != nil {
			client.Errorf("[scan end limit] Failed preparing table: %v", errors.Trace(err))
			return errors.Trace(err)
		}
	}

	// Create another table t2 with 3 rows (with out index), and it's value is long enough so that it will write to
	// TiKV's default cf.
	// Then these two tables are adjacent. If the iterator goes out of bound when scanning t1, it will touch t2's rows
	_, err = db.Exec("create table t2 (id integer, iv varchar(256));")
	if err != nil {
		client.Errorf("[scan end limit] Failed preparing table: %v", errors.Trace(err))
		return errors.Trace(err)
	}
	for i := 0; i < 3; i++ {
		_, err = db.Exec("insert into t2 values (?, ?);", i, randStr())
		if err != nil {
			client.Errorf("[scan end limit] Failed preparing table: %v", errors.Trace(err))
			return errors.Trace(err)
		}
	}

	// Get ids of those two tables
	table1Id, err := c.getTableId("t1")
	if err != nil {
		client.Errorf("[scan end limit] Failed getting table id for t1: %v", errors.Trace(err))
		return errors.Trace(err)
	}

	table2Id, err := c.getTableId("t2")
	if err != nil {
		client.Errorf("[scan end limit] Failed getting table id for t2: %v", errors.Trace(err))
		return errors.Trace(err)
	}

	client.Infof("[scan end limit] Created tables t1 (table id: %v) and t2 (table id: %v)", table1Id, table2Id)

	// Then delete default cf of rows of t2. So when t2's data is touched, TiKV will panic.
	pdAddr := c.cli.PdService()
	rawKvCli, err := tikv.NewRawKVClient([]string{pdAddr}, config.Security{})
	if err != nil {
		client.Errorf("[scan end limit] Failed connecting to RawKV: %v", errors.Trace(err))
		return errors.Trace(err)
	}
	defer rawKvCli.Close()

	scanStartKey := codec.EncodeBytes([]byte{}, tablecodec.EncodeTablePrefix(table2Id))
	client.Infof("[scan end limit] Delete default cf from key %+q", scanStartKey)

	keys, _, err := rawKvCli.Scan(scanStartKey, []byte{}, 3)
	if err != nil {
		client.Errorf("[scan end limit] Failed scanning raw keys: %v", errors.Trace(err))
		return errors.Trace(err)
	}

	for _, k := range keys {
		err = rawKvCli.Delete(k)
		if err != nil {
			client.Errorf("[scan end limit] Failed deleting raw key: %v", errors.Trace(err))
			return errors.Trace(err)
		}
	}

	return nil
}

// Runs the tests. Returns if test was canceled, test failed or test
func (c *ScanEndLimitCase) Run(ctx context.Context) error {
	dbAddr := c.cli.TiDBService()
	dbDSN := fmt.Sprintf("root:@tcp(%s)/%s", dbAddr, *dbName)
	db, err := util.OpenDB(dbDSN, *concurrency)
	if err != nil {
		client.Errorf("[scan end limit] OpenDB error %v", errors.Trace(err))
		return errors.Trace(err)
	}
	defer db.Close()

	c.startTime = time.Now()
	client.Infof("[scan end limit] Start running. Time limit: %v", c.timeLimit.String())

	// Start some threads to do normal CURD queries.
	var wg sync.WaitGroup
	var stopChannels []chan int
	for i := 0; i < c.queryThreads; i++ {
		wg.Add(1)
		stop := make(chan int)
		stopChannels = append(stopChannels, stop)
		go c.runQueriesWithDB(i, db, &wg, stop)
	}

	// Do operations that use kv_scan interface of TiKV.
	executedCommands := 0
	lastLogTime := time.Now()

LOOP:
	for {
		select {
		case <-ctx.Done():
			break LOOP
		default:
		}

		// Stop if timeLimit exceeded
		if c.timeLimit > 0 {
			if time.Since(c.startTime) >= c.timeLimit {
				break
			}
		}

		client.Infof("[scan end limit] admin check table t1...")
		_, err = db.Exec("admin check table t1;")
		if err != nil {
			client.Errorf("[scan end limit] Failed executing admin check table: %v", errors.Trace(err))
			break
		}

		client.Infof("[scan end limit] alter table t1 add index i1(i)...")
		_, err = db.Exec("alter table t1 add index i1(i);")
		if err != nil {
			client.Errorf("[scan end limit] Failed executing add index: %v", errors.Trace(err))
			break
		}

		client.Infof("[scan end limit] admin check table t1...")
		_, err = db.Exec("admin check table t1;")
		if err != nil {
			client.Errorf("[scan end limit] Failed executing admin check table: %v", errors.Trace(err))
			break
		}

		client.Infof("[scan end limit] alter table t1 drop index i1...")
		_, err = db.Exec("alter table t1 drop index i1;")
		if err != nil {
			client.Errorf("[scan end limit] Failed executing add index: %v", errors.Trace(err))
			break
		}

		executedCommands += 4
		// Periodically log amount of executed queries.
		if time.Since(lastLogTime) > c.logInterval {
			lastLogTime = time.Now()
			client.Infof("[scan end limit] executed %v critical commands")
		}
	}

	// Send stop signal to query threads.
	for _, channel := range stopChannels {
		channel <- 0
	}
	// Wait for queries threads finish.
	wg.Wait()

	if err != nil {
		client.Infof("[scan end limit] Test stopped with error")
	} else {
		client.Infof("[scan end limit] Test successfully finished")
	}

	return err
}

// Runs normal CURD queries on table `t1`.
func (c *ScanEndLimitCase) runQueriesWithDB(threadId int, db *sql.DB, wg *sync.WaitGroup, stop chan int) {
	defer wg.Done()

	executedQueries := 0
	lastLogTime := time.Now()

	for {
		// Stop if receives stop signal
		select {
		case <-stop:
			client.Infof("[scan end limit] thread %v has stopped", threadId)
			return
		default:
		}

		operation := rand.Int() % 5

		var err error
		switch operation {
		case 0:
			_, err = db.Exec("select * from t1 where id = ?;", rand.Int()%256)

		case 1:
			_, err = db.Exec("insert into t1 values (?, ?, ?) on duplicate key update i = ?, v = ?;", rand.Int()%256, rand.Int()%256, randStr(), rand.Int()%256, randStr())

		case 2:
			_, err = db.Exec("delete from t1 where id = ?;", rand.Int()%256)

		case 3:
			_, err = db.Exec("select sum(i) from t1;")

		case 4:
			_, err = db.Exec("update t1 set v = ? where i = ?;", randStr(), rand.Int()%256)
		}

		if err != nil {
			client.Errorf("[scan end limit] Failed executing query, err: %v", errors.Trace(err))
		}

		executedQueries++
		// Periodically log amount of executed queries.
		if time.Since(lastLogTime) > c.logInterval {
			lastLogTime = time.Now()
			client.Infof("[scan end limit] thread %v: executed %v queries", threadId, executedQueries)
		}
	}
}

// Gets table ID of the table with specified table name.
func (c *ScanEndLimitCase) getTableId(table string) (int64, error) {
	// TODO: Schrodinger should have an API to get TiDB status addr directly. Use that API when it's available.
	dbAddr := c.cli.TiDBService()
	dbStatusAddr := strings.Split(dbAddr, ":")[0] + ":10080"
	url := fmt.Sprintf("http://%s/schema/%s/%s", dbStatusAddr, c.dbName, table)

	resp, err := http.Get(url)
	if err != nil {
		return 0, errors.Trace(err)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return 0, errors.Trace(err)
	}

	if resp.StatusCode != 200 {
		return 0, errors.Errorf("HTTP request to TiDB status reporter returns %v. Body: %v", resp.StatusCode, string(body))
	}

	var data map[string]interface{}
	err = json.Unmarshal(body, &data)
	if err != nil {
		return 0, errors.Trace(err)
	}

	id := int64(data["id"].(float64))
	return id, nil
}

// Randomly generates a string of length 256.
func randStr() string {
	res := ""
	for i := 0; i < 256; i++ {
		res += string('a' + (rand.Int() % 26))
	}
	return res
}