package auth

import (
	"context"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/grafana/regexp"
	"github.com/pingcap/log"
	"github.com/tidbcloud/dataflow-service/app/container"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	accountAuth "github.com/tidbcloud/tiinf/sdk/auth"
	"go.uber.org/zap"
	"strconv"
)

func SetAuthOrgIdWithTiDBId(c *gin.Context) error {
	re := regexp.MustCompile(`tidbs/(\d+)`)
	matches := re.FindStringSubmatch(c.Request.URL.Path)

	if len(matches) > 1 {
		TiDBId, err := strconv.ParseUint(matches[1], 10, 64)
		if err != nil {
			return err
		}
		resp, err := container.Global().MgmtSvcGrpcClient.ListClusters(context.Background(), &mgmt.ListClustersRequest{Filter: &mgmt.ClusterFilter{ClusterIds: []uint64{TiDBId}}})
		if err != nil {
			return err
		}
		if len(resp.Clusters) == 0 {
			return fmt.Errorf("no cluster")
		}
		orgId := resp.Clusters[0].GetOrganizationId()
		query := c.Request.URL.Query()
		query.Set(accountAuth.AccessOrgIdQueryKey, strconv.FormatUint(orgId, 10))
		c.Request.URL.RawQuery = query.Encode()
		log.L().Info("set auth orgId with TiDBId", zap.Uint64("orgId", orgId), zap.Uint64("TiDBId", TiDBId))
	}
	return nil
}