package pkg

import (
	"bufio"
	"context"
	"database/sql"
	"embed"
	"fmt"
	"io"
	"io/fs"
	"io/ioutil"
	"math"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/onsi/gomega"
	"github.com/pingcap/endless/testcase/oltp/oltppkg"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"go.uber.org/zap"

	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/endless/pkg/workloadnode/tiup"
	"github.com/prometheus/client_golang/api"
)

func Sleep(ctx context.Context, timeout time.Duration) (isCtxDone bool) {
	select {
	case <-time.After(timeout):
		return false
	case <-ctx.Done():
		return true
	}
}

// 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
}

func AssignCmdStdout(cmd *workloadnode.Cmd) *io.PipeReader {
	pr, pw := io.Pipe()
	cmd.Stdout = pw
	return pr
}

func InfoCmdStdout(ctx context.Context, pr *io.PipeReader, errorHandler func(error) error) error {
	b := bufio.NewReader(pr)
	defer func() {
		_, err := ioutil.ReadAll(b)
		if err != nil {
			log.Warn("info cmd read all error", zap.Error(err))
		}
	}()
	for {
		select {
		case <-ctx.Done():
			log.Info("cmd stdout context done")
			return nil
		default:
			line, err := b.ReadString('\n')
			if err != nil {
				if err == io.EOF {
					log.Info("cmd stdout read eof")
					return nil
				}
				log.Error("read cmd stdout error", zap.String("line", line), zap.Error(err))
				return err
			}
			log.Info(line)
			if strings.Contains(strings.ToLower(line), "error") ||
				strings.Contains(strings.ToLower(line), "fail") ||
				strings.Contains(strings.ToLower(line), "fatal") {
				err = errors.New(line)
				if errorHandler != nil && errorHandler(err) != nil {
					log.Error("cmd stdout error", zap.Error(err))
					return err
				}
			}
		}
	}
}

func TiFlashExecCmd(tc resource.TiDBCluster, cmd string, peer int) (string, error) {
	log.Info("Start execute command on TiFlash nodes", zap.String("command", cmd), zap.Int("node", peer))
	stdout, stderr, exitCode, err := tc.Exec(resource.TiDBClusterExecOptions{Target: resource.TiFlash, Command: cmd, Peer: peer})
	if err != nil || exitCode != 0 {
		log.Error("Execute command error on TiFlash nodes",
			zap.String("command", cmd),
			zap.Int("node", peer),
			zap.Int("exit code", exitCode),
			zap.String("stderr", stderr),
			zap.String("stdout", stdout),
			zap.Error(err))
		return stdout, err
	}
	log.Info("End execute command on TiFlash nodes", zap.String("command", cmd), zap.Int("node", peer), zap.String("stdout", stdout))
	return stdout, nil
}

var OlapConParams = []string{"tidb_isolation_read_engines=tiflash", "tidb_mem_quota_query=32212254720", "multiStatements=true", "tidb_partition_prune_mode = 'dynamic'"}
var SqlConsistencyTiKVParma = []string{"tidb_isolation_read_engines=tikv", "multiStatements=true", "tidb_partition_prune_mode = 'dynamic'"}
var SqlConsistencyTiFlashParma = []string{"tidb_isolation_read_engines=tiflash", "multiStatements=true", "tidb_partition_prune_mode = 'dynamic'", "tidb_enforce_mpp=on"}

var MultiStatementsParam = "multiStatements=true"
var AllowAllFilesParam = "allowAllFiles=true"
var DynamicPartitionPruneMode = "tidb_partition_prune_mode = 'dynamic'"

func OpenDB(tc resource.Addressable, dbName string, params ...string) (*sql.DB, error) {
	addr, err := tc.ServiceURL(resource.DBAddr)
	if err != nil {
		return nil, err
	}
	// some ap query and ddl may exceed the timeout
	// params = append(params, "timeout=2m")
	// params = append(params, "readTimeout=2m")
	// params = append(params, "writeTimeout=2m")

	db, err := sql.Open("mysql", resource.BuildDSN(addr, dbName, params...))
	if err != nil {
		return nil, err
	}
	if err := db.Ping(); err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(50)
	db.SetMaxIdleConns(50)
	return db, nil
}

// DisableExpressionPushdown delete particular expression from pushdown blacklist and reload all tidb server
// Example:
// exprName = ["date_add"],  exprStoreType = ["tiflash"]
// sql will execute : DELETE FROM mysql.expr_pushdown_blacklist WHERE name = `date_add` and store_type = `tiflash`;
func DisableExpressionPushdown(tc resource.TiDBCluster, exprName []string, exprStoreType []string) error {
	db, err := tc.GetDB("")
	if err != nil {
		return err
	}
	defer db.Close()
	for i := range exprName {
		s := fmt.Sprintf("DELETE FROM mysql.expr_pushdown_blacklist WHERE name = '%s' and store_type = '%s'", exprName[i], exprStoreType[i])
		_, err := db.Exec(s)
		if err != nil {
			return err
		}
	}

	urls, err := tc.PeerURLs(resource.DBAddr)
	if err != nil {
		return err
	}
	reloadsql := "ADMIN reload expr_pushdown_blacklist"
	for _, url := range urls {
		var password = ""
		password, _ = url.User.Password()
		// url.Host format is tc-tidb-0.tc-tidb-peer.testbed-htap-dsj45:4000
		tidbaddr := fmt.Sprintf("%s:%s@tcp(%s)/%s", url.User.Username(), password, url.Host, "mysql")
		tidb, err := sql.Open("mysql", tidbaddr)
		if err != nil {
			return err
		}

		_, err = tidb.Exec(reloadsql)
		if err != nil {
			return err
		}

		err = tidb.Close()
		if err != nil {
			return err
		}
	}
	return nil
}

func ReadFileNamesAndContentsFromEmbedFs(queryFS embed.FS, path string) (names []string, contents []string, err error) {
	var dirEntries []fs.DirEntry
	dirEntries, err = queryFS.ReadDir(path)
	if err != nil {
		return
	}
	for _, r := range dirEntries {
		var c []byte
		c, err = queryFS.ReadFile(path + "/" + r.Name())
		if err != nil {
			return
		}
		names = append(names, r.Name())
		contents = append(contents, string(c))
	}
	return
}

func ReadFileNamesFromDir(path string) (names []string, err error) {
	if err = filepath.Walk(path, func(testPath string, info os.FileInfo, err error) error {
		if info == nil || info.IsDir() {
			return nil
		}
		names = append(names, testPath)
		return nil
	}); err != nil {
		return
	}
	return
}

func ReadFileNamesAndContentsFromDir(path string) (names []string, contents []string, err error) {
	if err = filepath.Walk(path, func(testPath string, info os.FileInfo, err error) error {
		if info == nil || info.IsDir() {
			return nil
		}
		names = append(names, testPath)
		return nil
	}); err != nil {
		return
	}
	for _, fn := range names {
		var c []byte
		c, err = ioutil.ReadFile(fn)
		if err != nil {
			return
		}
		contents = append(contents, string(c))
	}
	return
}

func GetLightningCmd(ctx context.Context, tc resource.TiDBCluster, s3Storage string) *workloadnode.Cmd {
	dbURL := matcher.Try(tc.ServiceURL(resource.DBAddr)).(*url.URL)
	pdURL := regexp.MustCompile("https?://").ReplaceAllString(matcher.Try(tc.ServiceURL(resource.PDAddr)).(*url.URL).String(), "")
	lighiningCmd := fmt.Sprintf(`/tidb-lightning "-pd-urls" "%s" "-tidb-host" "%s" "-tidb-port" "4000" "-tidb-user" "root" "-tidb-password" "" "-backend" "local" "-sorted-kv-dir" "/tmp/sorted-kv-dir" "-d" "%s"`, pdURL, dbURL.Hostname(), s3Storage)
	cmd := workloadnode.CommandContext(ctx, lighiningCmd)
	return cmd
}

// brStorage
// case 1:
// br 要求备份的数据和待恢复的集群 collation 信息一致（都包含或都不包含 collation）
// Br>=6.0 的备份数据中会记录 collation 的信息，Br>=6.0 恢复的时候会检测备份数据中是否包含 collation 信息，如果没有则报错
// 在系统测试中，为了减少维护的数据集数量，我们通常使用 4.0 的备份数据，但是上述 collation 的兼容性问题会引起 br 的数据不兼容性
// 如 daily run 的时候，使用 master 的 br 恢复 4.0 的数据，由于 4.0 的备份数据没有包含 collation 信息，br 恢复会失败
// 为了解决上述问题，我们需要维护两套低版本（4.0)的数据，分别是带有 collation 和不带有 collation 的数据；同时在恢复的时候需要设置 --check-requirements=false
// 目前 --check-requirements 参数会检测 collation 信息的一致性，br 和 tidb 集群的版本一致性
// case 2:
// br 新 feature 支持恢复系统表(user、permission)，原先 br 导出的备份也包含系统表但不支持恢复。
// 而 case 使用老版本备份(v4.0)，同时配置check-requirements = false，使得新版本 br 会恢复系统表，从而导致系统表列字段与新版本的不匹配。

// 我们引入了两个参数，br_storage 代表没有 collation 的数据，br_storage_with_collation 代表有 collation 的数据
// 通过检测当前集群是否打开了 collation 来判定具体使用哪一个备份数据
func MustRestoreWithBR(ctx context.Context, tc resource.TiDBCluster, br resource.BR, brStorage, brStorageWithCollation string) {
	tidbCli := tidb.NewClient(matcher.Try(tc.GetDB("")).(*sql.DB))
	defer tidbCli.Close()
	newCollationEnabled := matcher.Try(tidbCli.IsNewCollationEnabled()).(bool)
	log.Info("cluster new collation info", zap.Bool("enabled", newCollationEnabled))
	storage := brStorage
	if newCollationEnabled {
		storage = brStorageWithCollation
	}
	pdURL := regexp.MustCompile("https?://").ReplaceAllString(matcher.Try(tc.ServiceURL(resource.PDAddr)).(*url.URL).String(), "")
	brCmd := fmt.Sprintf(`/br restore full -u "%s" -s "%s" --send-credentials-to-tikv=true --check-requirements=false --filter '*.*' --filter '!mysql.*'`, pdURL, storage)
	cmd := workloadnode.CommandContext(ctx, brCmd)
	defer cmd.Stop()
	gomega.Expect(cmd.Run(br)).Should(gomega.BeNil())
}

func TiUPMustRestoreWithBR(ctx context.Context, tc *tiup.TiUPCluster, tiupNode resource.WorkloadNode, brStorage, brStorageWithCollation, version string, tls bool) {
	tidbCli := tidb.NewClient(matcher.Try(tc.GetDB("")).(*sql.DB))
	defer tidbCli.Close()
	newCollationEnabled := matcher.Try(tidbCli.IsNewCollationEnabled()).(bool)
	log.Info("cluster new collation info", zap.Bool("enabled", newCollationEnabled))
	storage := brStorage
	if newCollationEnabled {
		storage = brStorageWithCollation
	}
	// 当前 pd url 的形式是 pd-peer.endless-htap-ch-tls-tps-840706-1-148:2379
	// {host}.{namespace}:{port}
	// 在开启 tls 的时候，如果直接使用带有 namespace 的地址，则会报错
	// Failed to get config: Get "https://pd-peer.endless-htap-ch-tls-tps-840706-1-148:2379/pd/api/v1/config": x509: certificate is valid for localhost, pd-peer, not pd-peer.endless-htap-ch-tls-tps-840706-1-148
	// 有一些方法可以绕过上述错误
	// 1. tiup deploy 前，渲染 deploy.yaml 中 pd 的 host，加上 namespace 信息
	// 2. 使用 {host}:{port} 的形式。如果出现跨 namespace 访问 pd 的时候（如 case 访问资源中的 pd 链接），则访问不了
	// 简单起见，选择第二种方式
	pdURL := matcher.Try(tc.ServiceURL(resource.PDAddr)).(*url.URL)
	pdHost := strings.Split(regexp.MustCompile("https?://").ReplaceAllString(pdURL.Host, ""), ".")[0]
	brCmd := fmt.Sprintf(`tiup br:%s restore full -u "%s:%s" -s "%s" --send-credentials-to-tikv=true --check-requirements=false --filter '*.*' --filter '!mysql.*'`, version, pdHost, pdURL.Port(), storage)
	// todo use cluster name variable to replace "debug"
	if tls {
		brCmd += fmt.Sprint(` --ca="/root/.tiup/storage/cluster/clusters/debug/tls/ca.crt" --cert="/root/.tiup/storage/cluster/clusters/debug/tls/client.crt" --key="/root/.tiup/storage/cluster/clusters/debug/tls/client.pem"`)
	}
	cmd := workloadnode.CommandContext(ctx, brCmd)
	defer cmd.Stop()
	gomega.Expect(cmd.Run(tiupNode)).Should(gomega.BeNil())
}

func CheckTiFlashTiKVStoreSize(ctx context.Context, tc resource.TiDBCluster) error {
	log.Info("compare tiflash store compression ratio with tikv for single replica")
	tikvUsedSizes, err := oltppkg.GetNodeStoreSizeByType(tc, "tikv", "used_size")
	if err != nil {
		return err
	}
	db := matcher.Try(tc.GetDB("")).(*sql.DB)
	rows, err := db.Query("show config where type='pd' and name='replication.max-replicas'")
	if err != nil {
		return err
	}
	var tmp1, tmp2, tmp3, value string
	if rows.Next() {
		if err := rows.Scan(&tmp1, &tmp2, &tmp3, &value); err != nil {
			return err
		}
	} else {
		return errors.New("no tikv replica num")
	}
	tikvReplica, err := strconv.ParseFloat(value, 64)
	if err != nil {
		return err
	}
	var tikvStore float64
	for _, v := range tikvUsedSizes {
		tikvStore = tikvStore + v
	}
	tikvStore = tikvStore / tikvReplica

	tiflashUsedSizes, err := oltppkg.GetNodeStoreSizeByType(tc, "tiflash", "used_size")
	if err != nil {
		return err
	}
	rows, err = db.Query("select max(REPLICA_COUNT) from information_schema.tiflash_replica")
	if err != nil {
		return err
	}
	if rows.Next() {
		if err := rows.Scan(&value); err != nil {
			return err
		}
	} else {
		return errors.New("no tiflash replica num")
	}
	tiflashReplica, err := strconv.ParseFloat(value, 64)
	if err != nil {
		return err
	}
	var tiflashStore float64
	for _, v := range tiflashUsedSizes {
		tiflashStore = tiflashStore + v
	}
	tiflashStore = tiflashStore / tiflashReplica
	log.Info("tiflash store compression ratio", zap.Float64s("tikv store size", tikvUsedSizes), zap.Float64s("tiflash store size", tiflashUsedSizes), zap.Float64("tikv replica", tikvReplica), zap.Float64("tiflash replica", tiflashReplica), zap.Float64("tikv single replica", tikvStore), zap.Float64("tiflash single replica", tiflashStore), zap.Float64("compression ratio", tiflashStore/tikvStore))
	// TODO: compression ratio(1.5) should depends on workload, which should be parameter
	if tiflashStore/tikvStore > 1.5 {
		return errors.New("tiflash store compression ratio is more than 1.5 times")
	} else {
		return nil
	}
}

type ClusterInfo struct {
	Branch      string
	Version     string
	TidbHash    string
	PdHash      string
	TikvHash    string
	TiflashHash string
}

func PrintTiFlashVersionInformation(tiflash resource.WorkloadNode, clusterInfo *ClusterInfo) error {
	tiflashVersionCmd := "/tiflash/tiflash version"
	var info []string
	info, err := SyncRunCmdAndGetRes(tiflash, tiflashVersionCmd, "stdout")
	if err != nil {
		return err
	}
	for _, s := range info {
		if strings.HasPrefix(s, "Release Version:") {
			clusterInfo.Version = strings.Trim(s[16:], " ")
			log.Info("tiflash version:", zap.String("version", clusterInfo.Version))
		}
		if strings.HasPrefix(s, "Git Commit Hash:") {
			clusterInfo.TiflashHash = strings.Trim(s[16:], " ")
			log.Info("tiflash commit hash:", zap.String("hash", clusterInfo.TiflashHash))
			// Since cmd `/tiflash/tiflash version` will also print Raft proxy information, which have the same prefix = "Git Commit Hash:"). So break early.
			break
		}
	}
	return nil
}

// SyncRunCmdAndGetRes run command and put the stdout to a string array
// TODO : async for the method of log print
func SyncRunCmdAndGetRes(wn resource.WorkloadNode, cmd string, signal string) ([]string, error) {
	log.Info("Start execute command on TiFlash nodes", zap.String("command", cmd))
	stdout, stderr, exitCode, err := wn.Exec(resource.WorkloadNodeExecOptions{Command: cmd})
	if err != nil || exitCode != 0 {
		return nil, errors.New(fmt.Sprintf("Execute command error on TiFlash nodes. [command=%s][exitcode=%d][stderr=%s][stdout=%s]", cmd, exitCode, stderr, stdout))
	}
	var stdinfo []string
	// since TiFlash print log to different pipe, such as logs via stderr、version information via stdout
	if signal == "stdout" {
		stdinfo = strings.Split(stdout, "\n")
	} else if signal == "stderr" {
		stdinfo = strings.Split(stderr, "\n")
	}
	return stdinfo, nil
}

func GetMetric(promURL *url.URL, query string, ts time.Time) ([]float64, error) {
	prometheusAddr := promURL.String()
	client, err := api.NewClient(api.Config{
		Address: prometheusAddr,
	})
	if err != nil {
		return nil, err
	}
	httpAPI := v1.NewAPI(client)
	val, _, err := httpAPI.Query(context.Background(), query, ts)
	if err != nil {
		return nil, err
	}
	valMatrix := val.(model.Vector)
	if len(valMatrix) == 0 {
		return nil, errors.New("get empty metric")
	}
	var value []float64
	for i := range valMatrix {
		value = append(value, float64(valMatrix[i].Value))
		// judge weather exceeded float maximum value
		if math.IsNaN(value[i]) {
			return nil, errors.New(fmt.Sprintf("prometheus query result exceeded float maximum value, result=%s", valMatrix[i].String()))
		}
	}
	return value, nil
}