package host

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/pingcap/test-infra/caselib/pkg/model/workload"
	logsearch "github.com/pingcap/test-infra/sdk/pkg/logs"
	"github.com/pingcap/test-infra/sdk/resource"
	"gopkg.in/yaml.v2"
)

type TiUPClusterInfo struct {
	ClusterMeta TiUPClusterMeta `json:"cluster_meta,omitempty"`
	Instances   []TiUPInstance  `json:"instances,omitempty"`
}

type TiUPClusterMeta struct {
	ClusterType    string `json:"cluster_type"`
	ClusterName    string `json:"cluster_name"`
	ClusterVersion string `json:"cluster_version"`
	DeployUser     string `json:"deploy_user"`
	SSHType        string `json:"ssh_type"`
	TLSEnabled     bool   `json:"tls_enabled"`
	TLSCACert      string `json:"tls_ca_cert,omitempty"`
	TLSClientCert  string `json:"tls_client_cert,omitempty"`
	TLSClientKey   string `json:"tls_client_key,omitempty"`
	DashboardURL   string `json:"dashboard_url,omitempty"`
}

type TiUPInstance struct {
	ID        string `json:"id"`
	Role      string `json:"role"`
	Host      string `json:"host"`
	Ports     string `json:"ports"`
	OSArch    string `json:"os_arch"`
	Status    string `json:"status"`
	Since     string `json:"since"`
	DataDir   string `json:"data_dir"`
	DeployDir string `json:"deploy_dir"`

	ComponentName string
	Port          int
}

type TiUPCluster struct {
	diag *Diag

	haproxy             *haproxy
	ControlNodeHost     *WorkloadHost
	ControlNodeHostname string
	TestBed             string

	ClusterName string
	Topology    string
	CMDTimeout  time.Duration

	ClusterInfo TiUPClusterInfo
}

func NewTiUPCluster(controlNode resource.WorkloadNode, controlNodeHostname string, tiupTopo string, clusterName string) *TiUPCluster {
	// TODO: It's really dum to let user pass in controlNodeHostname
	// but there isn't method to get hostname of a workloadNode
	workloadHost := NewWorkloadHost(controlNode)
	tiupCluster := &TiUPCluster{
		TestBed:         controlNode.Handle().Testbed,
		ControlNodeHost: workloadHost,
		ClusterName:     clusterName,
		Topology:        tiupTopo,
		CMDTimeout:      10 * time.Minute,
		haproxy:         newHaproxy(workloadHost, controlNodeHostname),
	}
	tiupCluster.diag = newDiag(tiupCluster)

	return tiupCluster
}

// WithCMDTimeout sets timeout of tiup cluster operation
func (t *TiUPCluster) WithCMDTimeout(cmdTimeout time.Duration) *TiUPCluster {
	t.CMDTimeout = cmdTimeout
	return t
}

func (t *TiUPCluster) ServiceURL(addrType resource.TiDBAddrType) (*url.URL, error) {
	// except DBAddr, other addr types actually returns address of one of the instances
	switch addrType {
	case resource.DBAddr:
		return t.haproxy.getDBAddr()
	case resource.PDAddr:
		return t.getComponentAddr("pd")
	case resource.KVAddr:
		return t.getComponentAddr("tikv")
	case resource.TiFlashAddr:
		return t.getComponentAddr("tiflash")
	case resource.CDCAddr:
		return t.getComponentAddr("cdc")
	case resource.Prometheus:
		return t.getComponentAddr("prometheus")
	case resource.Alertmanager:
		return t.getComponentAddr("alertmanager")
	case resource.Grafana:
		return t.getComponentAddr("grafana")
	default:
		return nil, errors.New("service is unavailable")
	}
}

func (t *TiUPCluster) getComponentAddr(component string) (*url.URL, error) {
	clusterInfo, err := t.Display()
	if err != nil {
		return nil, err
	}

	for _, instance := range clusterInfo.Instances {
		if instance.ComponentName == component {
			return url.Parse(fmt.Sprintf("http://%s.%s:%d", instance.Host, t.TestBed, instance.Port))
		}
	}

	return nil, fmt.Errorf("component %s not found", component)
}

func (t *TiUPCluster) PeerURLs(addrType resource.TiDBAddrType) ([]*url.URL, error) {
	// except DBAddr, other addr types actually returns address of one of the instances
	switch addrType {
	case resource.DBAddr:
		val, err := t.haproxy.getDBAddr()
		if err != nil {
			return nil, err
		}
		return []*url.URL{val}, nil
	case resource.DBStatus:
		add, err := t.haproxy.getDBStatus()
		if err != nil {
			return nil, err
		}
		return []*url.URL{add}, nil
	case resource.PDAddr:
		return t.getAllComponentAddr("pd")
	case resource.KVAddr:
		return t.getAllComponentAddr("tikv")
	case resource.TiFlashAddr:
		return t.getAllComponentAddr("tiflash")
	case resource.CDCAddr:
		return t.getAllComponentAddr("cdc")
	case resource.Prometheus:
		return t.getAllComponentAddr("prometheus")
	case resource.Alertmanager:
		return t.getAllComponentAddr("alertmanager")
	case resource.Grafana:
		return t.getAllComponentAddr("grafana")
	default:
		return nil, errors.New("service is unavailable")
	}
}

func (t *TiUPCluster) getAllComponentAddr(component string) ([]*url.URL, error) {
	clusterInfo, err := t.Display()
	if err != nil {
		return nil, err
	}

	result := make([]*url.URL, 1)
	for _, instance := range clusterInfo.Instances {
		if instance.ComponentName == component {
			val, err := url.Parse(fmt.Sprintf("http://%s.%s:%d", instance.Host, t.TestBed, instance.Port))
			if err != nil {
				return nil, err
			}
			result = append(result, val)
		}
	}

	if len(result) > 0 {
		return result, nil
	}

	return nil, fmt.Errorf("component %s not found", component)
}

func (t *TiUPCluster) UpdateConfig(opts resource.TiDBClusterUpdateConfigOptions) error {
	clusterMetaConfig, err := t.getClusterMetaConfig()
	if err != nil {
		return err
	}

	configPatch := make(map[string]interface{})
	err = toml.Unmarshal([]byte(opts.Config), &configPatch)
	if err != nil {
		return err
	}
	configPatch = flattenMap(configPatch) // flatten configPatch map to match format of clusterMetaConfig

	t.updateClusterMetaConfig(clusterMetaConfig, opts.Target.String(), configPatch)
	newClusterMeta, err := yaml.Marshal(clusterMetaConfig)
	if err != nil {
		return err
	}

	// update meta.yaml on control node
	err = t.ControlNodeHost.workload.WriteFile(resource.WorkloadNodeWriteFileOptions{
		Content:  bytes.NewReader(newClusterMeta),
		FilePath: fmt.Sprintf("/root/.tiup/storage/cluster/clusters/%s/meta.yaml", t.ClusterName),
	})
	if err != nil {
		return err
	}

	// reload config to take effect
	return t.run(fmt.Sprintf("tiup cluster reload %s -R %s", t.ClusterName, opts.Target.String()), nil)
}

func (t *TiUPCluster) getClusterMetaConfig() (clusterMetaConfig map[string]interface{}, err error) {
	clusterMeta, err := t.ControlNodeHost.Exec(fmt.Sprintf("cat ~/.tiup/storage/cluster/clusters/%s/meta.yaml", t.ClusterName), t.CMDTimeout)
	if err != nil {
		return nil, err
	}

	clusterMetaConfig = make(map[string]interface{})
	err = yaml.Unmarshal([]byte(clusterMeta.Stdout), &clusterMetaConfig)
	if err != nil {
		return nil, err
	}

	return clusterMetaConfig, nil
}

// updateClusterMetaConfig inplace updates clusterMetaConfig with configPatch of component
func (t *TiUPCluster) updateClusterMetaConfig(clusterMetaConfig map[string]interface{}, component string, configPatch map[string]interface{}) {
	topology := clusterMetaConfig["topology"].(map[interface{}]interface{}) // topology must exist
	serverConfigs, ok := topology["server_configs"].(map[interface{}]interface{})
	if !ok {
		serverConfigs = make(map[interface{}]interface{})
	}

	// yaml.v2 treat every value as map[interface{}]interface, for reason we need to use patched version of v2(see comment in go.mod replace and PR)
	// so there is a type conversion: componentConfigMap => componentConfig before calling mergeMap()
	componentConfigMap, ok := serverConfigs[component].(map[interface{}]interface{})
	if !ok {
		componentConfigMap = make(map[interface{}]interface{})
	}
	componentConfig := make(map[string]interface{})
	for k, v := range componentConfigMap {
		componentConfig[k.(string)] = v
	}

	newComponentConfig := mergeMap(componentConfig, configPatch)
	serverConfigs[component] = newComponentConfig
	topology["server_configs"] = serverConfigs
	clusterMetaConfig["topology"] = topology
}

// Deploy will deploy tiup cluster with haproxy.
func (t *TiUPCluster) Deploy(version string, flags ...string) error {
	topoFilepath, err := writeTmpFile(t.ControlNodeHost.workload, "tiup.yaml", t.Topology)
	if err != nil {
		return err
	}

	err = t.run(fmt.Sprintf("tiup cluster deploy %s %s %s", t.ClusterName, version, topoFilepath), flags)
	if err != nil {
		return err
	}

	return t.deployHaproxy()
}

// deployHaproxy is idempotent, it terminates old haproxy process then starts a new one
func (t *TiUPCluster) deployHaproxy() error {
	clusterInfo, err := t.Display()
	if err != nil {
		return err
	}

	return t.haproxy.deploy(clusterInfo)
}

func (t *TiUPCluster) Start(flags ...string) error {
	return t.run(fmt.Sprintf("tiup cluster start %s", t.ClusterName), flags)
}

func (t *TiUPCluster) Stop(flags ...string) error {
	return t.run(fmt.Sprintf("tiup cluster stop %s", t.ClusterName), flags)
}

func (t *TiUPCluster) Restart(flags ...string) error {
	return t.run(fmt.Sprintf("tiup cluster restart %s", t.ClusterName), flags)
}

func (t *TiUPCluster) Destroy(flags ...string) error {
	err := t.haproxy.destroy()
	if err != nil {
		return err
	}

	return t.run(fmt.Sprintf("tiup cluster destroy %s", t.ClusterName), flags)
}

func (t *TiUPCluster) Upgrade(version string, flags ...string) error {
	return t.run(fmt.Sprintf("tiup cluster upgrade %s %s", t.ClusterName, version), flags)
}

func (t *TiUPCluster) ScaleIn(node string, flags ...string) error {
	err := t.run(fmt.Sprintf("tiup cluster scale-in %s -N %s", t.ClusterName, node), flags)
	if err != nil {
		return err
	}

	return t.deployHaproxy()
}

func (t *TiUPCluster) ScaleOut(tiupTopo string, flags ...string) error {
	topoFilepath, err := writeTmpFile(t.ControlNodeHost.workload, "tiup-scaleout.yaml", tiupTopo)
	if err != nil {
		return err
	}

	err = t.run(fmt.Sprintf("tiup cluster scale-out %s %s", t.ClusterName, topoFilepath), flags)
	if err != nil {
		return err
	}

	return t.deployHaproxy()
}

// Patch calls tiup cluster patch command with package from given packageURL
// package downloaded from packageURL will be re-packed to fulfill filepath requirement of tiup
//
// Explicit cluster stop is needed before starting an offline patch
func (t *TiUPCluster) Patch(packageURL string, downloadTimeout time.Duration, flags ...string) error {
	packageFilename := fmt.Sprintf("%s.tiup.tar.gz", time.Now().Format(time.RFC3339))
	packageFilepath := fmt.Sprintf("/tmp/%s", packageFilename)
	err := t.ControlNodeHost.Run(workload.NewCMD(
		fmt.Sprintf("curl -o %s %s", packageFilepath, packageURL), downloadTimeout))
	if err != nil {
		return err
	}

	packageFilepath, err = t.reTarPackage(packageFilepath)
	if err != nil {
		return err
	}

	return t.run(fmt.Sprintf("tiup cluster patch %s %s", t.ClusterName, packageFilepath), flags)
}

// reTarPackage unpacks oldPackagePath, mv files in ./bin/ to ./
// then re-pack all contents it into newPackagePath
//
// so that newPackagePath can be directly used in tiup cluster patch
func (t *TiUPCluster) reTarPackage(oldPackagePath string) (newPackagePath string, err error) {
	tmpPackageFolder := fmt.Sprintf("/tmp/%s", time.Now().Format(time.RFC3339))
	err = t.ControlNodeHost.Run(workload.NewCMD(fmt.Sprintf("mkdir %s", tmpPackageFolder), t.CMDTimeout))
	if err != nil {
		return "", err
	}

	err = t.ControlNodeHost.Run(workload.NewCMD(fmt.Sprintf("tar xzf %s --directory=%s", oldPackagePath, tmpPackageFolder), t.CMDTimeout))
	if err != nil {
		return "", err
	}

	// mv ./bin/* to ./
	// overwrite any conflict files
	err = t.ControlNodeHost.Run(workload.NewCMD(fmt.Sprintf(`
	DIR="%s"
	if [[ ! -e $DIR/bin ]]; then exit 0; fi;
	for file in $DIR/bin/*; do filename=$(basename $file); rm -rf $DIR/$filename; mv $DIR/bin/$filename $DIR/; done;`, tmpPackageFolder),
		t.CMDTimeout))
	if err != nil {
		return "", err
	}

	newPackagePath = fmt.Sprintf("/tmp/%s.tiup.tar.gz", time.Now().Format(time.RFC3339))
	err = t.ControlNodeHost.Run(workload.NewCMD(fmt.Sprintf("tar czf %s --directory=%s .", newPackagePath, tmpPackageFolder), t.CMDTimeout))
	if err != nil {
		return "", err
	}

	return newPackagePath, nil
}

// Display syncs TiUPCluster info with tiup cluster display $cluster-name
//
// Besides returning TiUPClusterInfo, the inner ClusterInfo field will also be set
func (t *TiUPCluster) Display() (TiUPClusterInfo, error) {
	var clusterInfo TiUPClusterInfo
	output, err := t.ControlNodeHost.Exec(fmt.Sprintf("tiup cluster display %s --uptime --format json", t.ClusterName), t.CMDTimeout)
	if err != nil {
		return clusterInfo, err
	}

	err = json.Unmarshal([]byte(output.Stdout), &clusterInfo)
	if err != nil {
		return clusterInfo, err
	}

	t.ClusterInfo = clusterInfo
	return clusterInfo, nil
}

func (t *TiUPCluster) run(cmd string, flags []string) error {
	flags = append(flags, "-y") // Skip all confirmations
	cmd += " "
	cmd += strings.Join(flags, " ")
	return t.ControlNodeHost.Run(workload.NewCMD(cmd, t.CMDTimeout))
}

func (t *TiUPCluster) GetDB(dbName string) (*sql.DB, error) {
	addr, err := t.ServiceURL(resource.DBAddr)
	if err != nil {
		return nil, err
	}

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

	if err = db.Ping(); err != nil {
		return nil, err
	}
	return db, nil
}

func (t *TiUPCluster) SearchLog(ctx context.Context, opts resource.TiDBClusterSearchLogOptions) ([]string, error) {
	var urls []*url.URL
	var err error
	switch opts.Target {
	case resource.TiDB:
		urls, err = t.PeerURLs(resource.DBStatus)
	case resource.PD:
		urls, err = t.PeerURLs(resource.PDAddr)
	case resource.TiKV:
		urls, err = t.PeerURLs(resource.KVAddr)
	default:
		return nil, fmt.Errorf("unsupported component for log search: %s", opts.Target.String())
	}
	if err != nil {
		return nil, err
	}

	logsClient := logsearch.NewDiagnosticLogClient()
	var logs []string
	for _, url := range urls {
		port, err := strconv.ParseInt(url.Port(), 10, 64)
		if err != nil {
			return nil, err
		}
		l, err := logsClient.SearchLog(ctx, url.Hostname(), int(port), opts.StartTime, opts.EndTime, opts.Levels, opts.Patterns, opts.Limit)
		if err != nil {
			return nil, err
		}
		logs = append(logs, l...)
	}
	return logs, nil
}

type haproxy struct {
	workloadHost *WorkloadHost
	hostName     string
	testbed      string

	dbPort       int
	dbStatusPort int
	config       string
}

func newHaproxy(host *WorkloadHost, hostName string) *haproxy {
	return &haproxy{workloadHost: host, hostName: hostName, testbed: host.workload.Handle().Testbed}
}

func (h *haproxy) getDBAddr() (*url.URL, error) {
	return url.Parse(fmt.Sprintf("mysql://root:@%s.%s:%d", h.hostName, h.testbed, h.dbPort))
}

func (h *haproxy) getDBStatus() (*url.URL, error) {
	return url.Parse(fmt.Sprintf("http://%s.%s:%d", h.hostName, h.testbed, h.dbStatusPort))
}

func (h *haproxy) deploy(clusterInfo TiUPClusterInfo) error {
	h.dbPort = defaultDBPort
	h.dbStatusPort = defaultDBStatusPort
	h.config = h.renderDefaultConfig(clusterInfo)

	configFilepath, err := writeTmpFile(h.workloadHost.workload, "haproxy.cfg", h.config)
	if err != nil {
		return err
	}

	// -sf terminates old pids
	return h.workloadHost.Run(
		workload.NewCMD(fmt.Sprintf("haproxy -f %s -p /run/haproxy.pid -sf $(cat /run/haproxy.pid)", configFilepath), time.Minute))
}

func (h *haproxy) destroy() error {
	return h.workloadHost.Run(workload.NewCMD("pkill haproxy", time.Minute))
}

func (h *haproxy) renderDefaultConfig(clusterInfo TiUPClusterInfo) string {

	dbPortHaConfig := `
listen tidb-cluster                        # Database load balancing.
   bind 0.0.0.0:3390                       # The Floating IP address and listening port.
   mode tcp                                # HAProxy uses layer 4, the transport layer.
   balance leastconn                       # The server with the smallest number of connections receives the connection. "leastconn" is recommended where long sessions are expected, such as LDAP, SQL and TSE, rather than protocols using short sessions, such as HTTP. The algorithm is dynamic, which means that server weights might be adjusted on the fly for slow starts for instance.
	`

	dbStatusPortHaConfig := `
listen tidb-cluster-status                        # Database load balancing.
   bind 0.0.0.0:10080                       # The Floating IP address and listening port.
   mode tcp                                # HAProxy uses layer 4, the transport layer.
   balance leastconn                       # The server with the smallest number of connections receives the connection. "leastconn" is recommended where long sessions are expected, such as LDAP, SQL and TSE, rather than protocols using short sessions, such as HTTP. The algorithm is dynamic, which means that server weights might be adjusted on the fly for slow starts for instance.
	`

	serverConfigTemplate := "\n   server %s %s:%d check inter 2000 rise 2 fall 3"
	for _, instance := range clusterInfo.Instances {
		if instance.ComponentName != "tidb" {
			continue
		}

		statusPort, _ := strconv.Atoi(strings.Split(instance.Ports, "/")[1])

		dbPortHaConfig += fmt.Sprintf(serverConfigTemplate, instance.Host, instance.Host, instance.Port)
		dbStatusPortHaConfig += fmt.Sprintf(serverConfigTemplate, instance.Host, instance.Host, statusPort)
	}

	return defaultHaproxyConfig + dbPortHaConfig + dbStatusPortHaConfig
}

const (
	defaultDBPort       = 3390 // value match defaultHaproxyConfig
	defaultDBStatusPort = 10080

	// from: https://docs.pingcap.com/tidb/stable/haproxy-best-practices
	defaultHaproxyConfig = `
global                                     # Global configuration.
   chroot      /var/lib/haproxy            # Changes the current directory and sets superuser privileges for the startup process to improve security.
   pidfile     /var/run/haproxy.pid        # Writes the PIDs of HAProxy processes into this file.
   maxconn     8000                        # The maximum number of concurrent connections for a single HAProxy process.
   user        haproxy                     # Same with the UID parameter.
   group       haproxy                     # Same with the GID parameter. A dedicated user group is recommended.
   nbproc      64                          # The number of processes created when going daemon. When starting multiple processes to forward requests, make sure the value is large enough so that HAProxy does not block processes.
   daemon                                  # Makes the process fork into background. It is equivalent to the command line "-D" argument. It can be disabled by the command line "-db" argument.
   stats socket /var/lib/haproxy/stats     # The directory where statistics output is saved.

defaults                                   # Default configuration.
   log global                              # Inherits the settings of the global configuration.
   retries 2                               # The maximum number of retries to connect to an upstream server. If the number of connection attempts exceeds the value, the backend server is considered unavailable.
   timeout connect  2s                     # The maximum time to wait for a connection attempt to a backend server to succeed. It should be set to a shorter time if the server is located on the same LAN as HAProxy.
   timeout client 30000s                   # The maximum inactivity time on the client side.
   timeout server 30000s                   # The maximum inactivity time on the server side.

`
)

func writeTmpFile(dstWorkloadNode resource.WorkloadNode, fileSuffix string, content string) (filepath string, err error) {
	filename := fmt.Sprintf("%s.%s", time.Now().Format(time.RFC3339), fileSuffix)
	filepath = fmt.Sprintf("/tmp/%s", filename)
	err = dstWorkloadNode.WriteFile(resource.WorkloadNodeWriteFileOptions{
		Content:  strings.NewReader(content),
		FilePath: filepath,
	})

	return filepath, err
}

// mergeMap recursively merge 'base' and 'patch' map
// value in 'patch' will overrides value in 'base' on conflict
func mergeMap(base map[string]interface{}, patch map[string]interface{}) map[string]interface{} {
	// copy from https://github.com/pingcap/test-infra/blob/7e778582c3510760ba12e1497fe9c993932d371c/sdk/resource/impl/k8s/tidbcluster_client.go#L433-L461
	result := base
	for key, value := range patch {
		old, exist := result[key]
		if !exist {
			result[key] = value
			continue
		}

		oldMap, oldOk := old.(map[string]interface{})
		newMap, mergeOk := value.(map[string]interface{})
		if oldOk && mergeOk {
			result[key] = mergeMap(oldMap, newMap)
		} else if result[key] != value {
			result[key] = value
		}
	}
	return result
}

// flattenMap convert mutil-layer map to single layer
func flattenMap(ms map[string]interface{}) map[string]interface{} {
	// before flattenMap()
	//   a:
	//     b:
	//       a: 3
	//       b: 2
	// after flattenMap()
	// a.b.a: 3
	// b.b.b: 2

	// copy from https: //github.com/pingcap/tiup/blob/52f7c868e7ba5e2c587ce309d9e8e6b1b42ac5db/pkg/cluster/spec/server_config.go#L118-L144
	result := map[string]interface{}{}
	for k, v := range ms {
		var sub map[string]interface{}

		if m, ok := v.(map[string]interface{}); ok {
			sub = flattenMap(m)
		} else if m, ok := v.(map[interface{}]interface{}); ok {
			fixM := map[string]interface{}{}
			for k, v := range m {
				if sk, ok := k.(string); ok {
					fixM[sk] = v
				}
			}
			sub = flattenMap(fixM)
		} else {
			result[k] = v
			continue
		}

		for sk, sv := range sub {
			result[k+"."+sk] = sv
		}
	}
	return result
}