package router

import (
	"os"
	"time"

	pbm "github.com/tidbcloud/aws-shared-provider/pbgen/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/metrics"
	"go.uber.org/zap"
)

func (r *Router) isCancelled() bool {
	select {
	case <-r.grpcCtx.Done():
		return true
	default:
		return false
	}
}

func (r *Router) startWatch() {
	r.watchWaitGroup.Add(1)
	go func() {
		r.logger.Info("backend watcher started")
		defer r.watchWaitGroup.Done()

		for {
			start := time.Now()
			stream, err := r.grpcClient.WatchBackendV2(r.grpcCtx, &pbm.WatchBackendRequestV2{
				Token: r.token, GatewayHostname: os.Getenv("HOSTNAME"),
			})
			if err != nil {
				if r.isCancelled() {
					return
				}
				r.logger.Error("WatchTidbBackendV2 error, try after 2 seconds", zap.Error(err))
				time.Sleep(2 * time.Second) //nolint:mnd // not need
				continue
			}
			r.logger.Info("grpc stream initialized", zap.Duration("duration", time.Since(start)))

			for {
				resp, err := stream.Recv()
				if err != nil {
					if r.isCancelled() {
						return
					}
					r.logger.Error("WatchTidbBackendV2 stream recv error, try after 2 seconds", zap.Error(err))
					time.Sleep(2 * time.Second) //nolint:mnd // not need
					break
				}

				r.logger.Info("recv backend", zap.Any("backend", resp))
				metrics.BackendRecvCounter.WithLabelValues(r.clusterID, r.token).Inc()

				pods := make([]*Pod, 0)
				for _, be := range resp.Backends {
					pods = append(pods, &Pod{IP: be.Ip, Name: be.PodName, TidbPool: resp.TidbPool})
				}
				err = r.cache.set(&Backend{
					TidbPool:    resp.TidbPool,
					ModRevision: resp.ModRevision,
					Replicas:    int(resp.Replicas),
					Pods:        pods,
				})
				if err != nil {
					r.logger.Debug("watch set backend skipped", zap.String("err", err.Error()), zap.Any("backend", resp))
				} else {
					r.logger.Info("watch set backend success", zap.Any("backend", resp))
				}

				// We can not refresh backend in zero backend since it will trigger reassign.
				// Refresh actively when there are one or more backends, it will not influence pod assignation.
				// If there are zero backend, the router will be refreshed passively by its own check health ticker.
				if len(resp.Backends) > 0 {
					r.RefreshBackend()
				}
			}
		}
	}()
}