package main

import (
	"database/sql"
	"flag"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"runtime"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/mux"
	"github.com/pingcap/errors"
	"github.com/pingcap/tidb-test/util"
	log "github.com/sirupsen/logrus"
)

var (
	addr      = flag.String("addr", ":9090", "http listen addr")
	tidbAddrs = flag.String("tidb-addr", "127.0.0.1:5001,127.0.0.1:5002,127.0.0.1:5003", "tidb addrs")
	pdAddrs   = flag.String("pd-addr", "127.0.0.1:42379", "pd addrs")
)

const (
	passResult = "pass"
)

var (
	keyGlobal int64 = 1
	servAddrs []string
)

func main() {
	flag.Parse()
	runtime.GOMAXPROCS(runtime.NumCPU())

	log.Debug("start probe")
	servAddrs = strings.Split(*tidbAddrs, ",")
	if len(servAddrs) != 3 {
		log.Fatal("bad arguments ", servAddrs)
	}
	runHTTPProbeResult()
}

func makeDBClient(addr string) *sql.DB {
	dsn := fmt.Sprintf("root@tcp(%s)/test?strict=true", addr)
	db, err := util.OpenDBWithRetry("mysql", dsn)
	if err != nil {
		log.Fatal(err)
	}
	log.Info("url:", dsn)

	return db
}

func runHTTPProbeResult() {
	m := mux.NewRouter()
	m.HandleFunc("/probe/server/start/leader", probeStartLeader)
	m.HandleFunc("/probe/server/start/follower", probeStartFollower)

	http.Handle("/", m)
	if err := http.ListenAndServe(*addr, nil); err != nil {
		log.Fatal("listen failed, err:", err)
	}
}

func getCurrentKey() string {
	return fmt.Sprintf("%08d", atomic.LoadInt64(&keyGlobal))
}

func generateKey() string {
	return fmt.Sprintf("%08d", atomic.AddInt64(&keyGlobal, 1))
}

func checkResult(err error, result string) error {
	if (err == nil && result == passResult) || (err != nil && result != passResult) {
		return nil
	}
	if err == nil && result != passResult {
		return errors.New("bad request")
	}

	return err
}

func probePass(isPass string, doOp func() error) error {
	var err error
	resultCh := make(chan error, 1)

	go func() {
		if err = doOp(); err != nil {
			resultCh <- err
		}
		resultCh <- nil
	}()

	timeout := time.After(5 * time.Second)
	select {
	case ret := <-resultCh:
		return checkResult(ret, isPass)
	case <-timeout:
		log.Warning("timeout:", isPass)
		return errors.New("timeout")
	}

	return nil
}

func getValue(r *http.Request, key, defaultVal string) string {
	val := r.FormValue(key)
	log.Info("val", val)
	if val == "" {
		val = defaultVal
	}

	return val
}

func putProbe(servID int, ret, tag string, count int) error {
	var err error
	if count < 0 {
		count = int(rand.Int63n(100) + 1)
	}

	DB := makeDBClient(servAddrs[servID])
	_, err = DB.Exec("drop table if exists t")
	if err != nil {
		return err
	}
	defer DB.Close()
	_, err = DB.Exec("create table t (a varchar(256))")
	if err != nil {
		return err
	}
	for i := 0; i < count; i++ {
		key := generateKey()
		err = probePass(ret,
			func() error {
				_, err1 := DB.Exec("insert into t values(?)", tag+key)
				return err1
			})
		if err != nil {
			break
		}
	}

	log.Info("end: put, keys:", count)
	return err
}

func probeOwner(servID int, isOwner bool) error {
	DB := makeDBClient(servAddrs[servID])
	ret, err := DB.Query("admin show ddl")
	if err != nil {
		return err
	}
	defer DB.Close()
	ret.Next()
	var ddlOwnerID, bgOwnerID, selfID string
	var a, b, c, d string
	err = ret.Scan(&a, &ddlOwnerID, &b, &c, &bgOwnerID, &d, &selfID)
	if err != nil {
		return err
	}
	defer ret.Close()

	isDDLOwner := ddlOwnerID == selfID
	isBgOwner := bgOwnerID == selfID
	msg := fmt.Sprintf("expect is owner %v, get ddl owner %v, bg owner %v", isOwner, isDDLOwner, isBgOwner)
	if isOwner {
		if !isDDLOwner || !isBgOwner {
			return errors.New(msg)
		}
		return nil
	}
	if isDDLOwner || isBgOwner {
		return errors.New(msg)
	}
	return nil
}

func probeStart(w http.ResponseWriter, r *http.Request, isOwner bool) {
	ret := getValue(r, "result", "pass")
	timeout := r.FormValue("timeout")
	t, err := strconv.Atoi(timeout)
	if err != nil {
		RespHTTPErr(w, http.StatusBadRequest, err.Error())
	}
	time.Sleep(time.Duration(t) * time.Second)
	idStr := r.FormValue("serv_id")
	servID, err := strconv.Atoi(idStr)
	if err != nil {
		RespHTTPErr(w, http.StatusBadRequest, err.Error())
	}
	log.Infof("probe start result %v, serv_id %v", ret, servID)

	if err := putProbe(servID, ret, "start", 5); err != nil {
		RespHTTPErr(w, http.StatusInternalServerError, err.Error())
		return
	}
	if err := probeOwner(servID, isOwner); err != nil {
		RespHTTPErr(w, http.StatusInternalServerError, err.Error())
		return
	}
	w.WriteHeader(http.StatusOK)
}

func RespHTTPErr(w http.ResponseWriter, code int, msg string) {
	w.WriteHeader(code)
	if msg == "" {
		msg = http.StatusText(code)
	}
	io.WriteString(w, msg)
}

func probeStartLeader(w http.ResponseWriter, r *http.Request) {
	log.Debug("start probe leader")
	probeStart(w, r, true)
	log.Debug("end probe leader")
}

func probeStartFollower(w http.ResponseWriter, r *http.Request) {
	log.Debug("start probe follower")
	probeStart(w, r, false)
	log.Debug("end probe follower")
}