package updatetest

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"testing"
	"time"

	_ "github.com/go-sql-driver/mysql"
	. "github.com/pingcap/check"
)

var (
	fromTidbCommitID = flag.String("oldtidb", "", "TiDB Commit ID")
	fromTikvCommitID = flag.String("oldtikv", "", "TiKV Commit ID")
	fromPdCommitID   = flag.String("oldpd", "", "PD Commit ID")
	toTidbCommitID   = flag.String("newtidb", "", "TiDB Commit ID")
	toTikvCommitID   = flag.String("newtikv", "", "TiKV Commit ID")
	toPdCommitID     = flag.String("newpd", "", "PD Commit ID")
)

const (
	tidbLogPath       = "log/tidb.log"
	tidbStderrLogPath = "log/tidb_stderr.log"
	tikvLogPath       = "log/tikv.log"
	tikvStderrLogPath = "log/tikv_stderr.log"
	pdLogPath         = "log/pd.log"
	pdStderrLogPath   = "log/pd_stderr.log"
	pdPath            = "127.0.0.1:2379"
)

func TestUpdate(t *testing.T) {
	CustomVerboseFlag = true
	TestingT(t)
}

var _ = Suite(&TestUpdateSuite{})

type TestUpdateSuite struct {
	tidbServer [2]*exec.Cmd
	tikvServer *exec.Cmd
	pdServer   *exec.Cmd
}

func init() {
	flag.Parse()
}

func (s *TestUpdateSuite) SetUpSuite(c *C) {

}

func (s *TestUpdateSuite) TestRollingUpdate(c *C) {
	if *fromPdCommitID != "" {
		getPDbyCommitID(*fromPdCommitID)
	}
	if *fromTidbCommitID != "" {
		getTiDBbyCommitID(*fromTidbCommitID)
	}
	if *fromTikvCommitID != "" {
		getTiKVbyCommitID(*fromTikvCommitID)
	}
	err := s.startPD()
	c.Assert(err, IsNil)
	err = s.startTiKV()
	c.Assert(err, IsNil)
	err = s.startTiDB(0)
	c.Assert(err, IsNil)
	err = s.startTiDB(1)
	c.Assert(err, IsNil)
	if *toPdCommitID != "" {
		go getPDbyCommitID(*toPdCommitID)
	}
	if *toTidbCommitID != "" {
		go getTiDBbyCommitID(*toTidbCommitID)
	}
	if *toTikvCommitID != "" {
		go getTiKVbyCommitID(*toTikvCommitID)
	}
	time.Sleep(20 * time.Second)
	err = s.killServer(s.pdServer.Process)
	c.Assert(err, IsNil)
	err = s.startPD()
	c.Assert(err, IsNil)
	err = s.killServer(s.tikvServer.Process)
	c.Assert(err, IsNil)
	err = s.startTiKV()
	c.Assert(err, IsNil)
	err = s.killServer(s.tidbServer[0].Process)
	c.Assert(err, IsNil)
	err = s.startTiDB(0)
	c.Assert(err, IsNil)
	err = s.killServer(s.tidbServer[1].Process)
	err = s.startTiDB(1)
	c.Assert(err, IsNil)
	addr0 := fmt.Sprintf("127.0.0.1:4000")
	addr1 := fmt.Sprintf("127.0.0.1:4001")
	for retry := 0; retry < 10; retry++ {
		var db1, db2 *sql.DB
		db1, _ = sql.Open("mysql", fmt.Sprintf("root@(%s)/test", addr0))
		db2, _ = sql.Open("mysql", fmt.Sprintf("root@(%s)/test", addr1))
		err1 := db1.Ping()
		err2 := db2.Ping()
		if err1 != nil || err2 != nil {
			time.Sleep(10 * time.Second)
		} else {
			break
		}
	}
	var db *sql.DB
	db, err = sql.Open("mysql", fmt.Sprintf("root@(%s)/test", addr0))
	c.Assert(err, IsNil)
	err = db.Ping()
	c.Assert(err, IsNil)
	db.Close()
	db, err = sql.Open("mysql", fmt.Sprintf("root@(%s)/test", addr1))
	c.Assert(err, IsNil)
	err = db.Ping()
	c.Assert(err, IsNil)
	err = s.killServer(s.tidbServer[1].Process)
	c.Assert(err, IsNil)
	err = s.killServer(s.tidbServer[0].Process)
	c.Assert(err, IsNil)
	err = s.killServer(s.tikvServer.Process)
	c.Assert(err, IsNil)
	err = s.killServer(s.pdServer.Process)
	c.Assert(err, IsNil)
}

func getTiDBbyCommitID(commitID string) {
	var url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/tidb/pr/" + commitID + "/centos7/tidb-server.tar.gz"
	cmd := exec.Command("curl", "-SLO", url)
	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
	cmd = exec.Command("tar", "zxvf", "tidb-server.tar.gz")
	if err := cmd.Run(); err != nil {
		log.Fatal("Get TiDB binary from PR failed !")
	}
}

func getTiKVbyCommitID(commitID string) {
	var url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/tikv/pr/" + commitID + "/unportable_centos7/tidb-server.tar.gz"
	cmd := exec.Command("curl", "-SLO", url)
	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
	cmd = exec.Command("tar", "zxvf", "tikv-server.tar.gz")
	if err := cmd.Run(); err != nil {
		log.Fatal("Get TiKV binary from PR failed !")
	}
}

func getPDbyCommitID(commitID string) {
	var url = "http://pingcap-dev.hk.ufileos.com/builds/pingcap/pd/pr/" + commitID + "/centos7/pd-server.tar.gz"
	cmd := exec.Command("curl", "-SLO", url)
	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
	cmd = exec.Command("tar", "zxvf", "pd-server.tar.gz", "-C")
	if err := cmd.Run(); err != nil {
		log.Fatal("Get PD binary from PR failed !")
	}
}

func (s *TestUpdateSuite) startTiDB(num int) error {
	cmd := exec.Command("bin/tidb-server",
		"--store=tikv",
		fmt.Sprintf("-P=%d", 4000+num),
		fmt.Sprintf("-socket=%s", ""),
		fmt.Sprintf("-path=%s", pdPath),
		fmt.Sprintf("-status=%d", 10080+num),
		fmt.Sprintf("--log-file=log/tidb%d.log", num))
	err := cmd.Start()
	if err != nil {
		return err
	}
	time.Sleep(5 * time.Second)
	var db *sql.DB
	addr := fmt.Sprintf("127.0.0.1:%d", 4000+num)
	tryCount := 10
	for i := 0; i < tryCount; i++ {
		db, err = sql.Open("mysql", fmt.Sprintf("root@(%s)/test", addr))
		if err != nil {
			log.Println("open addr %v failed, retry count %d err %v", addr, i, err)
			continue
		}
		err = db.Ping()
		if err == nil {
			break
		}
		log.Println("ping addr %v failed, retry count %d err %v", addr, i, err)

		db.Close()
		time.Sleep(10 * time.Second)
	}
	s.tidbServer[num] = cmd
	return nil
}

func (s *TestUpdateSuite) startTiKV() error {
	cmd := exec.Command("bin/tikv-server",
		"--addr=0.0.0.0:20171",
		fmt.Sprintf("--pd=%s", pdPath),
		"--data-dir=data1/deploy1/data",
		"--advertise-addr=127.0.0.1:20171",
		"--log-file=log/tikv.log")
	err := cmd.Start()
	if err != nil {
		return err
	}
	s.tikvServer = cmd
	time.Sleep(5 * time.Second)
	return nil
}

func (s *TestUpdateSuite) startPD() error {
	cmd := exec.Command("bin/pd-server",
		"--data-dir=data1/deploy1/data.pd",
		"--log-file=log/pd.log")
	err := cmd.Start()
	if err != nil {
		return err
	}
	s.pdServer = cmd
	time.Sleep(5 * time.Second)
	return nil
}

func (s *TestUpdateSuite) killServer(proc *os.Process) error {
	// Make sure this tidb is killed, and it makes the next tidb that has the same port as this one start quickly.
	err := proc.Kill()
	if err != nil {
		log.Fatal("kill server failed err %v", err)
		return err
	}
	_, err = proc.Wait()
	if err != nil {
		log.Fatal("kill server, wait failed err %v", err)
		return err
	}

	time.Sleep(1 * time.Second)
	return nil
}