package router

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	tiproxycfg "github.com/pingcap/tiproxy/lib/config"
	tiproxyrouter "github.com/pingcap/tiproxy/pkg/manager/router"
	pbm "github.com/tidbcloud/aws-shared-provider/pbgen/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"go.uber.org/zap"
	ctrl "sigs.k8s.io/controller-runtime"
)

const (
	routerIdleTimeout = 180 // seconds
)

// Router wraps tiproxy's router, it will determine routering request to which backend.
// It caches manager's tidb assignation and use watch method to keep the cache fresh.
type Router struct {
	*tiproxyrouter.ScoreBasedRouter

	logger       *zap.Logger
	routerBundle *RouterBundle
	cache        *Cache
	token        string
	clusterID    string
	lastAccess   atomic.Int64

	grpcClient     pbm.TidbBackendClient
	grpcCtx        context.Context
	watchWaitGroup sync.WaitGroup
	watchCancel    context.CancelFunc
	quit           chan struct{}
}

func NewRouter(
	ctx context.Context,
	logger *zap.Logger,
	token string,
	clusterID string,
	rb *RouterBundle,
	grpcClient pbm.TidbBackendClient,
) *Router {
	newLogger := logger.With(
		zap.String("component", "tiproxy"),
		zap.String("token", token),
		zap.String("cluster_id", clusterID),
	)

	sbr := tiproxyrouter.NewScoreBasedRouter(
		newLogger,
		tiproxyrouter.WithPauseEnabled(),
		tiproxyrouter.WithRebalanceInterval(time.Second),
	)

	childCtx, watchCancel := context.WithCancel(ctx)
	router := &Router{
		ScoreBasedRouter: sbr,
		logger:           newLogger,
		routerBundle:     rb,
		cache:            newCache(),
		token:            token,
		clusterID:        clusterID,
		grpcClient:       grpcClient,
		grpcCtx:          childCtx,
		watchCancel:      watchCancel,
		quit:             make(chan struct{}),
	}
	router.lastAccess.Store(time.Now().Unix())
	return router
}

func (r *Router) InitRouter(fetcher tiproxyrouter.BackendFetcher, hc tiproxyrouter.HealthCheck) error {
	err := r.Init(fetcher, hc, &tiproxycfg.HealthCheck{})
	if err != nil {
		ctrl.Log.Error(err, "failed to start score based router", "token", r.token, "cluster_id", r.clusterID)
		return utils.NewInternalError(utils.ERR_CLUSTER_TEMPORARY_UNAVAILABLE, err)
	}
	r.startWatch()
	go r.pollBackendLoop()
	return nil
}

func (r *Router) UpdateLastAccess() {
	r.lastAccess.Store(time.Now().Unix())
}

func (r *Router) GetBackendIPs() []string {
	pods := r.cache.listAllPods()
	result := make([]string, 0, len(pods))
	for _, pod := range pods {
		result = append(result, pod.IP)
	}
	return result
}

func (r *Router) GetPoolBackend(tidbPool string) ([]string, int, bool) {
	pods, replicas, replicasKnown := r.cache.getPoolBackend(tidbPool)
	result := make([]string, 0, len(pods))
	for _, pod := range pods {
		result = append(result, pod.IP)
	}
	return result, replicas, replicasKnown
}

func (r *Router) IdleTooLong() bool {
	return r.ConnCount() == 0 && time.Now().Unix()-r.lastAccess.Load() > routerIdleTimeout
}

func (r *Router) Close() {
	r.routerBundle.removeRouterItem(r.token)
	r.logger.Info("removed from router bundle")

	r.ScoreBasedRouter.Close()
	r.logger.Info("tiproxy router closed")

	close(r.quit)
	r.watchCancel()
	r.watchWaitGroup.Wait()
	r.logger.Info("backend watcher stopped")

	metrics.ClearMetrics(r.clusterID)
}