package slprober

import (
	"context"
	"fmt"
	"time"

	"github.com/pingcap/log"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"go.uber.org/zap"
	v1 "k8s.io/api/core/v1"

	"github.com/tidbcloud/sli-reporter/slprober/handshake"
)

type tcResult struct {
	success bool
	tcName  string
}

// probeTiDBCluster probes a TiDBCluster(tc)
// it will probe tidb or tiproxy (if exists) service related to this tc
// it only return alive status if both tidb and tiproxy are alive,
// otherwise return dead status
func probeTiDBCluster(
	ctx context.Context,
	tc *v1alpha1.TidbCluster,
	tidbSvc *v1.Service,
	tiproxySvc *v1.Service,
	resultCh chan<- tcResult,
	handshakeCli handshake.Interface,
) {
	clusterID := getClusterID(tc)
	timeBegin := time.Now()
	enableTiproxy := tc.Spec.TiProxy != nil && tc.Spec.TiProxy.Replicas > 0 && tiproxySvc != nil
	logger := log.L().With(zap.String("clusterID", clusterID),
		zap.String("tcName", tc.Name), zap.Bool("enableTiproxy", enableTiproxy))
	logger.Info("begin probe tc")

	probeServiceCount := 1
	if enableTiproxy {
		probeServiceCount++
	}
	serviceResultCh := make(chan serviceResult, probeServiceCount)
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	go probeTiDBService(ctx,
		fmt.Sprintf("%s:%d", tidbSvc.Spec.ClusterIP, getTiDBServicePort(tidbSvc)),
		tidbSvc.Name,
		serviceResultCh,
		handshakeCli,
	)
	if enableTiproxy {
		go probeTiDBService(ctx,
			fmt.Sprintf("%s:%d", tiproxySvc.Spec.ClusterIP, getTiproxyServicePort(tiproxySvc)),
			tiproxySvc.Name,
			serviceResultCh,
			handshakeCli,
		)
	}

	resultCount := 0
	for {
		select {
		case <-ctx.Done():
			logger.Info("finish probe tc, context done",
				zap.String("elapsed", time.Since(timeBegin).String()),
			)
			return
		case result := <-serviceResultCh:
			resultCount++
			if !result.success {
				// once we get an dead result, we can cancel the other goroutines
				logger.Info("finish probe tc, detected as dead",
					zap.String("elapsed", time.Since(timeBegin).String()),
				)
				resultCh <- tcResult{
					success: false,
					tcName:  tc.Name,
				}
				return
			}
			if resultCount >= probeServiceCount {
				// all services are alive
				logger.Info("finish probe tc, detected as alive",
					zap.String("elapsed", time.Since(timeBegin).String()),
				)
				resultCh <- tcResult{
					success: true,
					tcName:  tc.Name,
				}
				return
			}
		}
	}
}

func getTiDBServicePort(svc *v1.Service) int {
	if svc == nil {
		return 4000
	}
	for _, port := range svc.Spec.Ports {
		if port.Name == "mysql-client" {
			return int(port.Port)
		}
	}
	return 4000
}

func getTiproxyServicePort(svc *v1.Service) int {
	if svc == nil {
		return 6000
	}
	for _, port := range svc.Spec.Ports {
		if port.Name == "tiproxy-sql" {
			return int(port.Port)
		}
	}
	return 6000
}