package main

import (
	"context"
	"crypto/tls"
	"sync"
	"time"

	"github.com/pingcap/log"
	"github.com/pingcap/tidb/pkg/util"
	"github.com/tidbcloud/aws-shared-provider/pkg/datacleaner/common"
	pdhttp "github.com/tidbcloud/aws-shared-provider/pkg/utils/pd"
	"github.com/tikv/client-go/v2/config"
	"github.com/tikv/client-go/v2/txnkv"
	pd "github.com/tikv/pd/client"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.uber.org/zap"
)

var (
	once          sync.Once
	globalClients *clients
)

func getGlobalClients() *clients {
	once.Do(func() {
		c, err := newPDAndEtcdGroupClients()
		if err != nil {
			log.Fatal("Failed to create clients", zap.Error(err))
		}
		globalClients = c
	})
	return globalClients
}

type clients struct {
	ctx          context.Context
	cancel       context.CancelFunc
	tls          *tls.Config
	PdClient     pd.Client
	HttpPdClient pdhttp.HttpPdClient
	TikvClient   *txnkv.Client
	PDEtcdClient *clientv3.Client

	// For meta service group (etcd group)
	MetaServiceGroupEtcdClient map[string]*clientv3.Client
}

func (c *clients) GetPdContext() *common.PdContext {
	return common.NewPdContext(c.tls, []string{*pdAddr})
}

func (c *clients) Close() error {
	c.cancel()
	c.PdClient.Close()
	err := c.TikvClient.Close()
	if err != nil {
		return err
	}
	err = c.PDEtcdClient.Close()
	if err != nil {
		return err
	}
	return nil
}

func newPDAndEtcdGroupClients() (c *clients, err error) {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	defer func() {
		if err != nil {
			cancel()
		}
	}()

	tlsConfig, err := util.NewTLSConfig(
		util.WithCAPath(*ca),
		util.WithCertAndKeyPath(*cert, *key),
	)
	if err != nil {
		return nil, err
	}

	// create pd client
	pdClient, err := pd.NewClientWithContext(ctx, []string{*pdAddr}, pd.SecurityOption{
		CAPath:   *ca,
		CertPath: *cert,
		KeyPath:  *key,
	}, pd.WithCustomTimeoutOption(time.Duration(*pdTimeout)*time.Second))
	if err != nil {
		return nil, err
	}

	// create http pd client
	httpPdClient := pdhttp.NewClient(*pdAddr, tlsConfig)

	// create tikv client
	config.UpdateGlobal(func(conf *config.Config) {
		conf.Security = config.Security{
			ClusterSSLCA:   *ca,
			ClusterSSLCert: *cert,
			ClusterSSLKey:  *key,
		}
		conf.TiKVClient.MaxBatchSize = 0
	})
	tikvClient, err := txnkv.NewClient([]string{*pdAddr})
	if err != nil {
		return nil, err
	}

	// create PD etcd clients
	pdEtcdCfg := clientv3.Config{
		Endpoints: []string{*pdAddr},
		TLS:       tlsConfig,
	}
	pdEtcdClient, err := clientv3.New(pdEtcdCfg)
	if err != nil {
		return nil, err
	}

	metaServiceGroupEtcdClient, err := getMetaServiceGroupEtcdClient(ctx, httpPdClient)
	if err != nil {
		return nil, err
	}

	return &clients{
		ctx:                        ctx,
		cancel:                     cancel,
		tls:                        tlsConfig,
		PdClient:                   pdClient,
		HttpPdClient:               httpPdClient,
		TikvClient:                 tikvClient,
		PDEtcdClient:               pdEtcdClient,
		MetaServiceGroupEtcdClient: metaServiceGroupEtcdClient,
	}, nil
}

func getMetaServiceGroupEtcdClient(ctx context.Context, httpPdClient pdhttp.HttpPdClient) (map[string]*clientv3.Client, error) {
	// pd http client for meta service group operations
	metaServiceGroupList, err := httpPdClient.GetMetaServiceGroupList(ctx)
	if err != nil {
		return nil, err
	}
	log.Info("Get meta service group from PD", zap.Any("metaServiceGroupList", metaServiceGroupList))

	metaServiceGroupEtcdClient := make(map[string]*clientv3.Client)
	for _, metaServiceGroup := range *metaServiceGroupList {
		etcdAddr := metaServiceGroup.Addresses
		if etcdAddr == "" {
			log.Warn("Meta service group etcd address is empty", zap.String("group", metaServiceGroup.Addresses))
			continue
		}

		// create etcd client for each group
		client, err := newEtcdClient(etcdAddr)
		if err != nil {
			log.Error("Failed to create etcd client for meta service group", zap.String("group", metaServiceGroup.Addresses), zap.Error(err))
			continue
		}
		metaServiceGroupEtcdClient[metaServiceGroup.ID] = client
		log.Info("Created etcd client for meta service group", zap.String("groupID", metaServiceGroup.ID), zap.Any("etcdAddr", client.Endpoints()))
	}
	return metaServiceGroupEtcdClient, nil
}

func newEtcdClient(etcdAddr string) (*clientv3.Client, error) {
	tlsConfig, err := util.NewTLSConfig(
		util.WithCAPath(*ca),
		util.WithCertAndKeyPath(*cert, *key),
	)
	if err != nil {
		return nil, err
	}
	// create etcd client
	etcdCfg := clientv3.Config{
		Endpoints: []string{etcdAddr},
		TLS:       tlsConfig,
	}
	return clientv3.New(etcdCfg)
}