package util

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/url"
	"os"
	"os/exec"
	"path"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql" // mysql driver
	"github.com/juju/errors"
	"github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/push"
)

const (
	alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"

	// Used by RandString
	letterIdxBits = 6                    // 6 bits to represent a letter index
	letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
	letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
)

// PushPrometheus pushes metrics to Prometheus Pushgateway.
func PushPrometheus(job, addr string, interval time.Duration) {
	for {
		err := push.New(addr, job).Gatherer(prometheus.DefaultGatherer).Push()
		if err != nil {
			client.Errorf("[%s] could not push metrics to Prometheus Pushgateway: %v", job, err)
		}

		time.Sleep(interval)
	}
}

// RandString reference: http://stackoverflow.com/questions/22892120/how-to-generate-a-random-string-of-a-fixed-length-in-golang
func RandString(b []byte, r *rand.Rand) {
	n := len(b)
	// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
	for i, cache, remain := n-1, r.Int63(), letterIdxMax; i >= 0; {
		if remain == 0 {
			cache, remain = r.Int63(), letterIdxMax
		}
		if idx := int(cache & letterIdxMask); idx < len(alphabet) {
			b[i] = alphabet[idx]
			i--
		}
		cache >>= letterIdxBits
		remain--
	}
}

// RunWithRetry tries to run func in specified count
func RunWithRetry(ctx context.Context, retryCnt int, interval time.Duration, f func() error) error {
	var (
		err error
	)
	for i := 0; retryCnt < 0 || i < retryCnt; i++ {
		err = f()
		if err == nil {
			return nil
		}

		select {
		case <-ctx.Done():
			return nil
		case <-time.After(interval):
		}
	}
	return errors.Trace(err)
}

// QueryEntry is a query
type QueryEntry struct {
	Query              string
	Args               []interface{}
	ExpectAffectedRows int64
}

// MustExec must execute sql or fatal
func MustExec(db *sql.DB, query string, args ...interface{}) sql.Result {
	r, err := db.Exec(query, args...)
	if err != nil {
		client.Fatalf("exec %s err %v", query, err)
	}
	return r
}

// OpenDB opens db
func OpenDB(dsn string, maxIdleConns int) (*sql.DB, error) {
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, err
	}

	db.SetMaxIdleConns(maxIdleConns)
	client.Info("DB opens successfully")
	return db, nil
}

// ExecWithRollback exeutes or rollbak
func ExecWithRollback(db *sql.DB, queries []QueryEntry) (res sql.Result, err error) {
	tx, err := db.Begin()
	if err != nil {
		return nil, errors.Trace(err)
	}
	for _, q := range queries {
		res, err = tx.Exec(q.Query, q.Args...)
		if err != nil {
			tx.Rollback()
			return nil, errors.Trace(err)
		}
		if q.ExpectAffectedRows >= 0 {
			affected, err := res.RowsAffected()
			if err != nil {
				tx.Rollback()
				return nil, errors.Trace(err)
			}
			if affected != q.ExpectAffectedRows {
				client.Fatalf("expect affectedRows %v, but got %v, query %v", q.ExpectAffectedRows, affected, q)
			}
		}
	}
	if err = tx.Commit(); err != nil {
		tx.Rollback()
		return nil, errors.Trace(err)
	}
	return
}

// IsFileExist returns true if the file exists.
func IsFileExist(name string) bool {
	_, err := os.Stat(name)
	return err == nil
}

// Wget downloads a string URL to the dest directory and returns the file path.
// SKips if the file already exists.
func Wget(ctx context.Context, rawURL string, dest string) (string, error) {
	u, err := url.Parse(rawURL)
	if err != nil {
		return "", err
	}

	if len(dest) == 0 {
		dest = "."
	}

	fileName := path.Base(u.Path)
	filePath := path.Join(dest, fileName)
	if IsFileExist(filePath) {
		return filePath, nil
	}

	os.MkdirAll(dest, 0755)
	err = exec.CommandContext(ctx, "wget", "--tries", "20", "--waitretry", "60",
		"--retry-connrefused", "--dns-timeout", "60", "--connect-timeout", "60",
		"--read-timeout", "60", "--directory-prefix", dest, rawURL).Run()
	return filePath, errors.Annotatef(err, "wget %s to %s", rawURL, dest)
}

// InstallArchive downloads the URL and extracts the archive to the dest diretory.
// Supports zip, and tarball.
func InstallArchive(ctx context.Context, rawURL string, dest string) error {
	var name string
	var err error
	if strings.HasPrefix(rawURL, "file://") {
		name = strings.Trim(rawURL, "file://")
	} else {
		if name, err = Wget(ctx, rawURL, "/tmp/chaos"); err != nil {
			return err
		}
	}

	if strings.HasSuffix(name, ".zip") {
		err = exec.CommandContext(ctx, "unzip", "-d", dest, name).Run()
	} else {
		err = exec.CommandContext(ctx, "tar", "-xf", name, "-C", dest).Run()
	}
	if err != nil {
		return errors.Annotate(err, "decompression error")
	}
	return nil
}

// ExecCmd execs cmd
func ExecCmd(ctx context.Context, command string) (string, error) {
	args := strings.Split(command, " ")
	name := args[0]
	args = args[1:]
	cmd := exec.CommandContext(ctx, name, args...)
	output, err := cmd.CombinedOutput()
	return string(output), err
}

// DropTable execs drop table
func DropTable(db *sql.DB, tableName string) error {
	stmt := fmt.Sprintf("DROP TABLE IF EXISTS `%s`", tableName)
	_, err := db.Exec(stmt)
	return errors.Trace(err)
}

// CreateTable execs create table
func CreateTable(db *sql.DB, tableName, tableSchema string) error {
	stmt := fmt.Sprintf("CREATE TABLE IF NOT EXISTS `%s` %s", tableName, tableSchema)
	_, err := db.Exec(stmt)
	return errors.Trace(err)
}

// RecreateTable drops table and then creates table
func RecreateTable(db *sql.DB, tableName, tableSchema string) error {
	if err := DropTable(db, tableName); err != nil {
		return err
	}
	return errors.Trace(CreateTable(db, tableName, tableSchema))
}

const (
	asciiStart = int('a')
	asciiLimit = int('z')
)

// RandomASCIIBytes generates bytes randomly
func RandomASCIIBytes(r *rand.Rand, size int) []byte {
	data := make([]byte, size)
	for i := range data {
		data[i] = byte(r.Intn(asciiLimit-asciiStart) + asciiStart)
	}
	return data
}

// ReadJSON reads json data from io.ReaderCloser
func ReadJSON(r io.ReadCloser, data interface{}) error {
	defer r.Close()

	b, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}

	return json.Unmarshal(b, data)
}

// DumpJSON dumps json data
func DumpJSON(obj interface{}, format bool) (string, error) {
	data, err := json.Marshal(obj)
	if err != nil {
		return "", err
	}
	if format {
		var out bytes.Buffer
		json.Indent(&out, data, "", "\t")
		data = out.Bytes()
	}
	return string(data), nil
}

// SliceToMap converts slice to map
func SliceToMap(slice []string) map[string]interface{} {
	sMap := make(map[string]interface{})
	for _, str := range slice {
		sMap[str] = struct{}{}
	}
	return sMap
}

// RandDistinct generates `count` distinct uniformly-distributed rand numbers within the range `[0, max)`
func RandDistinct(count, max int) []int {
	// replaceMap simulates an array sized `max` and with it's items equals to their indices.
	replaceMap := map[int]int{}
	res := make([]int, 0, count)

	replace := func(v int) int {
		replaced, ok := replaceMap[v]
		if ok {
			return replaced
		}
		return v
	}

	for i := 0; i < count; i++ {
		// Rand in [i, max)
		v := rand.Intn(max-i) + i
		// map[v] is the current result
		rv := replace(v)
		res = append(res, rv)
		// Swap map[i] and map[v], but since map[i] won't be used again, only map[v] need to be set.
		replaceMap[v] = replace(i)
	}
	return res
}

// TestTiDBReplicaRead checks `tidb_replica_read` variable can be set/get correctly.
func TestTiDBReplicaRead(job string, dbDSN string) {
	var db, db1 *sql.DB
	var err error

	db, err = OpenDB(dbDSN, 1)
	if err != nil {
		client.Fatalf("[%s] open %s fail: %v", job, dbDSN, err)
	}
	db1, err = OpenDB(dbDSN, 1)
	if err != nil {
		client.Fatalf("[%s] open %s fail: %v", job, dbDSN, err)
	}

	// Check default value of `tidb_replica_read`.
	replicaRead := getTiDBReplicaRead(job, db)
	if replicaRead != "leader" {
		client.Fatalf("[%s] tidb_replica_read should be leader by default", job)
	}

	// Check change `tidb_replica_read`.
	setTiDBReplicaRead(job, db, "follower")
	replicaRead = getTiDBReplicaRead(job, db)
	if replicaRead != "follower" {
		client.Fatalf("[%s] tidb_replica_read should be follower after set", job)
	}

	// `tidb_replica_read` of db1 should be still `leader`.
	replicaRead = getTiDBReplicaRead(job, db1)
	if replicaRead != "leader" {
		client.Fatalf("[%s] tidb_replica_read should be leader because in the session it's not changed", job)
	}
}

func getTiDBReplicaRead(job string, db *sql.DB) string {
	var replicaRead string
	if err := db.QueryRow("select @@tidb_replica_read").Scan(&replicaRead); err != nil {
		client.Fatalf("[%s] get tidb_replica_read fail: %v", job, err)
	}
	return replicaRead
}

func setTiDBReplicaRead(job string, db *sql.DB, mode string) {
	sql := fmt.Sprintf("set @@tidb_replica_read = '%s'", mode)
	if _, err := db.Exec(sql); err != nil {
		client.Fatalf("[%s] tidb_replica_read set fail: %v", job, err)
	}
}

// RandomlyChangeReplicaRead changes `tidb_replica_read` randomly in a session.
func RandomlyChangeReplicaRead(job, replicaRead string, db *sql.DB) {
	if replicaRead != "leader" && rand.Float32() < 0.3 {
		if getTiDBReplicaRead("bank2", db) == "leader" {
			setTiDBReplicaRead("bank2", db, replicaRead)
		} else {
			setTiDBReplicaRead("bank2", db, "leader")
		}
	}
}