package router

import (
	"context"
	"sync"

	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/client"
	"github.com/tidbcloud/aws-shared-provider/pkg/gateway/config"
	"github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"go.uber.org/zap"
)

// RounterBundle manages all active token's routers. It organizes routers as a thread safe map.
type RouterBundle struct {
	sync.RWMutex
	routers map[string]*Router

	regionalInfra infra.RegionalInfra
	infraApi      infra.InfraAPI
	localKube     infra.LocalKube
	httpClient    *client.ManagerServerClient
	tbClient      pbm.TidbBackendClient
}

func NewRouterBundle(
	regionalInfra infra.RegionalInfra,
	infraApi infra.InfraAPI,
	localKube infra.LocalKube,
	httpClient *client.ManagerServerClient,
	tbClient pbm.TidbBackendClient,
) *RouterBundle {
	return &RouterBundle{
		regionalInfra: regionalInfra,
		infraApi:      infraApi,
		localKube:     localKube,
		httpClient:    httpClient,
		tbClient:      tbClient,
		routers:       make(map[string]*Router),
	}
}

func (rb *RouterBundle) getRouterItem(token string) *Router {
	rb.RLock()
	defer rb.RUnlock()
	return rb.routers[token]
}

func (rb *RouterBundle) removeRouterItem(token string) {
	rb.Lock()
	defer rb.Unlock()
	delete(rb.routers, token)
}

func (rb *RouterBundle) GetBackendTidbPool(token, backend string) string {
	rt := rb.getRouterItem(token)
	if rt == nil {
		return ""
	}
	pods := rt.cache.listAllPods()
	for _, pod := range pods {
		if backend == pod.IP {
			return pod.TidbPool
		}
	}
	return ""
}

func (rb *RouterBundle) GetBackendIPs(token string) []string {
	rt := rb.getRouterItem(token)
	if rt == nil {
		return nil
	}

	return rt.GetBackendIPs()
}

// GetConnCountByToken returns the number of connections for the given token as maintained by
// tiproxy's router. This is the authoritative count (tiproxy updates it on connect/disconnect)
// and does not rely on our OnConnClose being called. Returns 0 if no router exists for the token.
func (rb *RouterBundle) GetConnCountByToken(token string) int {
	rt := rb.getRouterItem(token)
	if rt == nil {
		return 0
	}
	return rt.ConnCount()
}

// TryGetRouter tries to get exist router from router bundle. If not found, it returns nil.
func (rb *RouterBundle) TryGetRouter(token string) tiproxyrouter.Router {
	rb.RLock()
	router := rb.routers[token]
	rb.RUnlock()
	if router == nil {
		return nil
	}
	router.UpdateLastAccess()
	return router
}

// CreateRouter creates a new router and add it to router bundle.
func (rb *RouterBundle) CreateRouter(
	ctx context.Context,
	logger *zap.Logger,
	token string,
	clusterID string,
	conf *config.Config,
) (tiproxyrouter.Router, error) {
	router := NewRouter(ctx, logger, token, clusterID, rb, rb.tbClient)
	rb.Lock()
	if oldRouter := rb.routers[token]; oldRouter != nil {
		rb.Unlock()
		oldRouter.UpdateLastAccess()
		return oldRouter, nil
	}
	rb.routers[token] = router
	rb.Unlock()

	fetcher := &backendFetcher{
		token:         token,
		clusterID:     clusterID,
		conf:          conf,
		regionalInfra: rb.regionalInfra,
		infraApi:      rb.infraApi,
		httpClient:    rb.httpClient,
		router:        router,
	}
	hc := &healthChecker{
		router:    router,
		token:     token,
		clusterID: clusterID,
		localKube: rb.localKube,
	}

	err := router.InitRouter(fetcher, hc)
	if err != nil {
		return nil, err
	}

	return router, nil
}

// ConnCount returns the total number of connections in the router bundle.
func (rb *RouterBundle) ConnCount() int {
	rb.RLock()
	defer rb.RUnlock()

	result := 0
	for _, r := range rb.routers {
		result += r.ConnCount()
	}
	return result
}