package topsql

import (
	"bytes"
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	. "github.com/onsi/gomega"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/pkg/proxy"
	"github.com/pingcap/test-infra/sdk/resource"
	"github.com/pingcap/tidb-dashboard/pkg/utils/topology"
	"github.com/tidwall/gjson"
	"go.etcd.io/etcd/clientv3"
	"go.uber.org/zap"
	netproxy "golang.org/x/net/proxy"
	"google.golang.org/grpc"

	_ "github.com/pingcap/test-infra/sdk/resource/impl/k8s"

	"github.com/pingcap/endless/pkg/matcher/metrics"
	"github.com/pingcap/endless/pkg/workloadnode"
	"github.com/pingcap/endless/pkg/workloadnode/tiup"
)

const DefaultVersion = "nightly"

var tiupTopoTmp = `
global:
  user: "tidb-%s"
  deploy_dir: "/tiup-%s/deploy"
  data_dir: "/tiup-%s/data"

pd_servers:
  - host: pd-peer
    client_port: %d
    peer_port: %d

tidb_servers:
  - host: tidb-peer
    port: %d
    status_port: %d

tikv_servers:
  - host: tikv-peer
    port: %d
    status_port: %d
    config:
      # disable load base split to avoid conflict with topsql
      split.region-cpu-overload-threshold-ratio: 0.0

monitoring_servers:
  - host: tiup-peer
    port: %d
    %s

%s
`

type ClusterConfig struct {
	Version              string
	Name                 string
	PDClientPort         uint16
	PDPeerPort           uint16
	TiDBPort             uint16
	TiDBStatusPort       uint16
	TiKVPort             uint16
	TiKVStatusPort       uint16
	PrometheusPort       uint16
	NodeExporterPort     uint16
	BlackboxExporterPort uint16
	NgmPort              uint16
}

func (c *ClusterConfig) buildTopologyYaml() string {
	ngmCfg := ""
	if c.NgmPort > 0 {
		ngmCfg = fmt.Sprintf("ng_port: %d", c.NgmPort)
	} else {
		ngmCfg = "ng_port: -1"
	}
	monitoredCfg := ""
	if c.NodeExporterPort > 0 {
		monitoredCfg += "monitored:\n"
		monitoredCfg += fmt.Sprintf("  node_exporter_port: %d\n", c.NodeExporterPort)
	}
	if c.BlackboxExporterPort > 0 {
		if monitoredCfg == "" {
			monitoredCfg += "monitored:\n"
		}
		monitoredCfg += fmt.Sprintf("  blackbox_exporter_port: %d\n", c.BlackboxExporterPort)
	}

	return fmt.Sprintf(tiupTopoTmp,
		c.Name, c.Name, c.Name,
		c.PDClientPort, c.PDPeerPort,
		c.TiDBPort, c.TiDBStatusPort,
		c.TiKVPort, c.TiKVStatusPort,
		c.PrometheusPort, ngmCfg, monitoredCfg)
}

type ClusterBundle struct {
	cfg ClusterConfig

	ctx    context.Context
	cancel context.CancelFunc

	tiupNode resource.WorkloadNode
	Cluster  *tiup.TiUPCluster
	DB       *sql.DB

	pdPorts map[uint16]struct{}
	etcdCli *clientv3.Client

	sysbenchNode resource.WorkloadNode
	sysbenchCmd  map[uint16]*workloadnode.Cmd

	LatestTopSQLOffUnixTime int64
}

func NewClusterBundle(tiupNode, sysbenchNode resource.WorkloadNode, clusterCfg ClusterConfig) (b *ClusterBundle, err error) {
	if clusterCfg.Version == "" {
		clusterCfg.Version = DefaultVersion
	}
	ctx, cancel := context.WithCancel(context.Background())
	b = &ClusterBundle{
		cfg:          clusterCfg,
		ctx:          ctx,
		cancel:       cancel,
		tiupNode:     tiupNode,
		sysbenchNode: sysbenchNode,
		sysbenchCmd:  map[uint16]*workloadnode.Cmd{},
		pdPorts:      map[uint16]struct{}{clusterCfg.PDClientPort: {}},
		Cluster:      tiup.NewTiUPCluster(tiupNode, "tiup-peer", clusterCfg.buildTopologyYaml(), clusterCfg.Name),
	}
	if err = b.Cluster.Deploy(b.cfg.Version); err != nil {
		return
	}
	if err = b.Cluster.Start(); err != nil {
		return
	}

	b.DB, err = b.getDB("", clusterCfg.TiDBPort)
	if err != nil {
		return
	}
	b.DB.SetConnMaxLifetime(time.Minute * 3)
	b.DB.SetMaxOpenConns(10)
	b.DB.SetMaxIdleConns(10)

	err = b.initEtcdClient()
	return
}

func (b *ClusterBundle) TurnOnTopSQL() error {
	if b.DB == nil {
		return errors.New("DB is nil")
	}

	_, err := b.DB.ExecContext(b.ctx, "SET GLOBAL tidb_enable_top_sql = on")
	return err
}

func (b *ClusterBundle) TurnOffTopSQL() error {
	if b.DB == nil {
		return errors.New("DB is nil")
	}

	_, err := b.DB.ExecContext(b.ctx, "SET GLOBAL tidb_enable_top_sql = off")
	b.LatestTopSQLOffUnixTime = time.Now().Unix()
	return err
}

func (b *ClusterBundle) TopSQLIsOn() (res bool, err error) {
	if b.DB == nil {
		return false, errors.New("DB is nil")
	}

	row := b.DB.QueryRowContext(b.ctx, "SELECT @@GLOBAL.tidb_enable_top_sql")
	err = row.Scan(&res)
	return
}

func (b *ClusterBundle) NgmURL(ngPort uint16) string {
	return fmt.Sprintf("http://tiup-peer.%s:%d", b.Cluster.TestBed, ngPort)
}

func (b *ClusterBundle) NgmTopology() (string, error) {
	if b.etcdCli == nil {
		return "", errors.New("etcd client is nil")
	}
	return topology.FetchNgMonitoringTopology(b.ctx, b.etcdCli)
}

func (b *ClusterBundle) NgmConnections() ([]string, error) {
	return ngmConnections(b.tiupNode)
}

func (b *ClusterBundle) PrepareSysbench(tidbPort uint16) error {
	if b.DB == nil {
		return errors.New("DB is nil")
	}

	_, err := b.DB.ExecContext(b.ctx, "CREATE DATABASE sbtest")
	if err != nil {
		return err
	}

	return workloadnode.CommandContext(b.ctx, fmt.Sprintf(`/usr/bin/sysbench \
		--db-driver=mysql \
		--mysql-host=%s \
		--mysql-port=%d \
		--mysql-user=root \
		--tables=10 \
		--table-size=10000 \
		oltp_point_select \
		prepare`, "tidb-peer", tidbPort)).Run(b.sysbenchNode)
}

func (b *ClusterBundle) DestroySysbench() error {
	if b.DB == nil {
		return errors.New("DB is nil")
	}

	_, err := b.DB.ExecContext(b.ctx, "DROP DATABASE sbtest")
	return err
}

func (b *ClusterBundle) StartSysbench(tidbPort uint16) error {
	if b.DB == nil {
		return errors.New("DB is nil")
	}

	if _, ok := b.sysbenchCmd[tidbPort]; ok {
		return errors.New("sysbench is running")
	}

	cmd := workloadnode.CommandContext(b.ctx, fmt.Sprintf(`/usr/bin/sysbench \
		--db-driver=mysql \
		--mysql-host=%s \
		--mysql-port=%d \
		--mysql-user=root \
		--tables=10 \
		--table-size=10000 \
		--time=30000 \
		--threads=1 \
		oltp_point_select \
		run`, "tidb-peer", tidbPort))
	b.sysbenchCmd[tidbPort] = cmd
	return cmd.Start(b.sysbenchNode)
}

func (b *ClusterBundle) StopSysbench(tidbPort uint16) {
	cmd, ok := b.sysbenchCmd[tidbPort]
	if !ok {
		return
	}

	_ = cmd.Stop()
	_ = cmd.Wait()
	delete(b.sysbenchCmd, tidbPort)
}

func (b *ClusterBundle) StopAllSysbench() {
	for k, v := range b.sysbenchCmd {
		_ = v.Stop()
		_ = v.Wait()
		delete(b.sysbenchCmd, k)
	}
}

func (b *ClusterBundle) ScaleOut(component string, portFirst, portSecond uint16) error {
	scaleOutTmp := ""
	switch component {
	case "tidb":
		scaleOutTmp = `
tidb_servers:
  - host: tidb-peer
    port: %d
    status_port: %d
`
	case "tikv":
		scaleOutTmp = `
tikv_servers:
  - host: tikv-peer
    port: %d
    status_port: %d
    config:
      # disable load base split to avoid conflict with topsql
      split.region-cpu-overload-threshold-ratio: 0.0
`
	case "pd":
		scaleOutTmp = `
pd_servers:
  - host: pd-peer
    client_port: %d
    peer_port: %d
`
	case "monitoring":
		scaleOutTmp = `
monitoring_servers:
  - host: tiup-peer
    port: %d
    ng_port: %d
`
	default:
		return errors.New("component is valid for tidb, tikv, pd or monitoring")
	}

	yml := fmt.Sprintf(scaleOutTmp, portFirst, portSecond)
	if err := b.Cluster.ScaleOut(yml); err != nil {
		return err
	}

	if component == "pd" {
		b.pdPorts[portFirst] = struct{}{}
		return b.initEtcdClient()
	}
	return nil
}

func (b *ClusterBundle) ScaleIn(component string, port uint16) error {
	// some pre-processing
	switch component {
	case "pd":
		// evict leader
		if err := workloadnode.CommandContext(b.ctx, fmt.Sprintf(`
tiup ctl:%s pd -u http://pd-peer:%d member leader show | grep %d -q && \
tiup ctl:%s pd -u http://pd-peer:%d member leader resign || \
true
`, b.cfg.Version, port, port, b.cfg.Version, port)).Run(b.tiupNode); err != nil {
			return err
		}
	case "tikv":
		// set replicas to 1 to avoid failure
		for pdPort := range b.pdPorts {
			if err := workloadnode.CommandContext(b.ctx, fmt.Sprintf(`
tiup ctl:%s pd -u http://pd-peer:%d config set max-replicas 1 || true
`, b.cfg.Version, pdPort)).Run(b.tiupNode); err != nil {
				return err
			}
		}
	}

	node := ""
	if component == "monitoring" {
		node = fmt.Sprintf("tiup-peer:%d", port)
	} else {
		node = fmt.Sprintf("%s-peer:%d", component, port)
	}
	if err := b.Cluster.ScaleIn(node); err != nil {
		return err
	}

	if component == "pd" {
		delete(b.pdPorts, port)
		return b.initEtcdClient()
	}

	return nil
}

func (b *ClusterBundle) StopComponent(component string, port uint16) error {
	node := component
	if component == "monitoring" {
		node = "tiup"
	}
	return b.Cluster.Stop(fmt.Sprintf("-N %s-peer:%d", node, port))
}

func (b *ClusterBundle) RestartComponent(component string, port uint16) error {
	return b.Cluster.Restart(fmt.Sprintf("-N %s-peer:%d", component, port))
}

type Target struct {
	StatusAddr string
	PubSubAddr string
}

func (b *ClusterBundle) ExpectTopSQLRunning(g Gomega, ngmURL string, targets []Target) {
	now := time.Now().Unix()
	b.ExpectTargetsProfiling(g, targets)
	ExpectDataAvailable(g, ngmURL, targets, now-120, now)
}

func (b *ClusterBundle) ExpectTargetsProfiling(g Gomega, targets []Target) {
	for _, target := range targets {
		if strings.HasPrefix(target.StatusAddr, "tidb") {
			g.Expect(b.Cluster).Should(metrics.TiDBProfilingCPU(target.StatusAddr))
		} else if strings.HasPrefix(target.StatusAddr, "tikv") {
			g.Expect(b.Cluster).Should(metrics.TiKVRsmeterStatingThreads(target.StatusAddr))
		}
	}
}

func ExpectDataAvailable(g Gomega, ngmURL string, targets []Target, startTs, endTs int64) {
	instances, err := FetchInstances(ngmURL, startTs, endTs)
	g.Expect(err).ShouldNot(HaveOccurred())
	g.Expect(QueryJSON(instances, "status")).Should(Equal("ok"))

	for _, target := range targets {
		var instance string
		var instanceType string
		if strings.HasPrefix(target.PubSubAddr, "tidb") {
			instance = target.PubSubAddr
			instanceType = "tidb"
		} else if strings.HasPrefix(target.PubSubAddr, "tikv") {
			instance = target.PubSubAddr
			instanceType = "tikv"
		} else {
			continue
		}

		g.Expect(QueryJSON(instances, "data")).Should(ContainElement(map[string]interface{}{"instance": instance, "instance_type": instanceType}))

		summary, err := FetchSummary(ngmURL, startTs, endTs, instance, instanceType)
		g.Expect(err).ShouldNot(HaveOccurred())
		g.Expect(QueryJSON(summary, "status")).Should(Equal("ok"))
		g.Expect(QueryJSON(summary, "data")).Should(Not(BeEmpty()))
	}
}

func (b *ClusterBundle) Close() {
	_ = b.closeEtcd()
	b.StopAllSysbench()
	_ = b.DestroySysbench()
	_ = workloadnode.CommandContext(b.ctx, "pkill haproxy").Run(b.tiupNode)
	_ = workloadnode.CommandContext(b.ctx, fmt.Sprintf("tiup cluster destroy %s -y", b.cfg.Name)).Run(b.tiupNode)
	b.cancel()
}

func (b *ClusterBundle) getDB(dbName string, port uint16) (*sql.DB, error) {
	addr, err := url.Parse(fmt.Sprintf("http://tidb-peer.%s:%d", b.Cluster.TestBed, port))
	if err != nil {
		return nil, err
	}

	db, err := sql.Open("mysql", resource.BuildDSN(addr, dbName))
	if err != nil {
		return nil, err
	}

	if err = db.PingContext(b.ctx); err != nil {
		return nil, err
	}
	return db, nil
}

func (b *ClusterBundle) initEtcdClient() (err error) {
	if b.etcdCli != nil {
		if err = b.etcdCli.Close(); err != nil {
			return
		}
		b.etcdCli = nil
	}

	var endpoints []string
	for k := range b.pdPorts {
		endpoints = append(endpoints, fmt.Sprintf("pd-peer.%s:%d", b.Cluster.TestBed, k))
	}
	dialer := proxy.GrpcProxyDialer()
	zapCfg := zap.NewProductionConfig()
	zapCfg.Encoding = log.ZapEncodingName
	b.etcdCli, err = clientv3.New(clientv3.Config{
		Context:   b.ctx,
		Endpoints: endpoints,
		DialOptions: []grpc.DialOption{
			grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) {
				switch d := dialer.(type) {
				case netproxy.ContextDialer:
					return d.DialContext(ctx, "tcp", addr)
				default:
					return d.Dial("tcp", addr)
				}
			}),
		},
		LogConfig: &zapCfg,
	})

	return
}

func (b *ClusterBundle) closeEtcd() error {
	if b.etcdCli == nil {
		return nil
	}
	err := b.etcdCli.Close()
	b.etcdCli = nil
	return err
}

func (b *ClusterBundle) ngmConfigured() bool {
	return b.cfg.NgmPort > 0
}

type NgmBundle struct {
	ctx    context.Context
	cancel context.CancelFunc

	clusterName string
	ngmPort     uint16
	node        resource.WorkloadNode

	ngmURL *url.URL
	ngmCmd *workloadnode.Cmd

	LatestNgmDownUnixTime int64

	configPath string
	logPath    string
	dataPath   string
}

func NewNgmBundle(ngmNode, tiupNode resource.WorkloadNode, ngmPort, pdPort, prometheusPort uint16, clusterName string) (b *NgmBundle, err error) {
	if ngmPort == 0 {
		return nil, errors.New("ngm is not configured")
	}

	ctx, cancel := context.WithCancel(context.Background())
	b = &NgmBundle{
		ctx:         ctx,
		cancel:      cancel,
		clusterName: clusterName,
		ngmPort:     ngmPort,
		node:        ngmNode,
	}
	err = workloadnode.CommandContext(b.ctx, fmt.Sprintf("mkdir -p /tiup-%s/deploy/ngm-%d/bin", clusterName, ngmPort)).Run(ngmNode)
	if err != nil {
		return
	}
	err = workloadnode.CommandContext(b.ctx, fmt.Sprintf("mkdir -p /tiup-%s/data/ngm-%d", clusterName, ngmPort)).Run(ngmNode)
	if err != nil {
		return
	}
	b.configPath = fmt.Sprintf("/tiup-%s/data/ngm-%d/config.toml", clusterName, ngmPort)
	err = workloadnode.CommandContext(b.ctx, fmt.Sprintf(`echo -e '[pd]\nendpoints = ["pd-peer:%d"]' > %s`, pdPort, b.configPath)).Run(ngmNode)
	if err != nil {
		return
	}
	err = workloadnode.CommandContext(b.ctx, fmt.Sprintf(
		`scp -o "UserKnownHostsFile=/dev/null" -o "StrictHostKeyChecking=no" \
		/tiup-%s/deploy/prometheus-%d/bin/ng-monitoring-server \
		%s-peer:/tiup-%s/deploy/ngm-%d/bin/ng-monitoring-server`,
		clusterName, prometheusPort, ngmNode.Handle().Name, clusterName, ngmPort,
	)).Run(tiupNode)
	if err != nil {
		return
	}

	b.logPath = fmt.Sprintf("/tiup-%s/data/ngm-%d/log", clusterName, ngmPort)
	b.dataPath = fmt.Sprintf("/tiup-%s/data/ngm-%d/data", clusterName, ngmPort)

	b.ngmURL, err = url.Parse(fmt.Sprintf("http://%s-peer.%s:%d", ngmNode.Handle().Name, ngmNode.Handle().Testbed, ngmPort))
	if err != nil {
		return
	}
	return
}

func (b *NgmBundle) StartNGM() error {
	if b.ngmCmd != nil {
		return errors.New("ng-monitoring-server is running")
	}

	b.ngmCmd = workloadnode.CommandContext(b.ctx, fmt.Sprintf(`/tiup-%s/deploy/ngm-%d/bin/ng-monitoring-server \
		--address=%s-peer:%d \
		--advertise-address=%s-peer:%d \
		--log.path=%s \
		--storage.path=%s \
		--config=%s`,
		b.clusterName, b.ngmPort,
		b.node.Handle().Name, b.ngmPort,
		b.node.Handle().Name, b.ngmPort,
		b.logPath, b.dataPath, b.configPath))
	err := b.ngmCmd.Start(b.node)
	if err != nil {
		return err
	}

	// wait for ready
	isReady := func() bool {
		var resp *http.Response
		resp, err = http.Get(fmt.Sprintf("%s/health", b.NgmURL()))
		if err != nil {
			return false
		}
		err = resp.Body.Close()
		return err == nil
	}

	deadline := time.Now().Add(5 * time.Second)
	for !isReady() && time.Now().Before(deadline) {
		time.Sleep(500 * time.Millisecond)
	}
	return err
}

func (b *NgmBundle) StopNGM() error {
	if b.ngmCmd == nil {
		return nil
	}

	if err := b.ngmCmd.Stop(); err != nil {
		return err
	}
	if err := b.ngmCmd.Wait(); err != nil {
		return err
	}

	b.ngmCmd = nil
	b.LatestNgmDownUnixTime = time.Now().Unix()
	return nil
}

func (b *NgmBundle) NgmURL() string {
	return b.ngmURL.String()
}

func (b *NgmBundle) ChangePDEndpoint(pdPort uint16) (err error) {
	if b.ngmURL == nil {
		return errors.New("ng-monitoring-server isn't running")
	}

	err = workloadnode.CommandContext(b.ctx, fmt.Sprintf(`echo -e '[pd]\nendpoints = ["pd-peer:%d"]' > %s`, pdPort, b.configPath)).Run(b.node)
	if err != nil {
		return
	}

	return workloadnode.CommandContext(b.ctx, "kill -1 $(pgrep ng-monitoring)").Run(b.node)
}

func (b *NgmBundle) Connections() ([]string, error) {
	return ngmConnections(b.node)
}

func (b *NgmBundle) Close() {
	_ = b.StopNGM()
	_ = workloadnode.CommandContext(b.ctx, fmt.Sprintf("rm -rf /tiup-%s/deploy/ngm-%d", b.clusterName, b.ngmPort)).Run(b.node)
	_ = workloadnode.CommandContext(b.ctx, fmt.Sprintf("rm -rf /tiup-%s/data/ngm-%d", b.clusterName, b.ngmPort)).Run(b.node)
	b.cancel()

	b.node = nil
	b.ngmCmd = nil
	b.ngmURL = nil
	b.configPath = ""
	b.logPath = ""
	b.dataPath = ""
}

func FetchInstances(ngmURL string, startTs, endTs int64) ([]byte, error) {
	resp, err := http.Get(fmt.Sprintf("%s/topsql/v1/instances?start=%d&end=%d", ngmURL, startTs, endTs))
	if err != nil {
		return nil, err
	}
	return ioutil.ReadAll(resp.Body)
}

func FetchSummary(ngmURL string, startTs, endTs int64, instance, instanceType string) ([]byte, error) {
	step := (endTs - startTs) / 300
	if step == 0 {
		step = 1
	}

	resp, err := http.Get(
		fmt.Sprintf("%s/topsql/v1/summary?start=%d&end=%d&step=%ds&instance=%s&instance_type=%s&top=5",
			ngmURL,
			startTs, endTs, step,
			instance, instanceType,
		))
	if err != nil {
		return nil, err
	}
	return ioutil.ReadAll(resp.Body)
}

func TurnOffConprof(ngmURL string) error {
	resp, err := http.Post(fmt.Sprintf("%s/config", ngmURL), "application/json", bytes.NewBuffer([]byte(`{"continuous_profiling":{"enable":false}}`)))
	if err != nil {
		return err
	}
	return resp.Body.Close()
}

func ngmConnections(node resource.WorkloadNode) ([]string, error) {
	stdout, stderr, exitCode, err := node.Exec(resource.WorkloadNodeExecOptions{Command: `ss -tHpn | grep $(pgrep ng-monitor) | awk '{print $5}' | sort | uniq`})
	if exitCode != 0 {
		return nil, errors.New(stderr)
	}
	if err != nil {
		return nil, err
	}

	return strings.Fields(stdout), nil
}

func CheckTiUPAvailable(node resource.WorkloadNode) error {
	_, stderr, exitCode, err := node.Exec(resource.WorkloadNodeExecOptions{Command: "tiup update --self && tiup list"})
	if exitCode != 0 {
		return errors.New(stderr)
	}
	if err != nil {
		return err
	}
	return nil
}

func QueryJSON(bytes []byte, path string) interface{} {
	return gjson.GetBytes(bytes, path).Value()
}