package adapterimpl

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"time"

	"github.com/go-resty/resty/v2"
	"github.com/google/uuid"
	"golang.org/x/sync/singleflight"

	"github.com/tidbcloud/tidb-management-service/serverless-global-service/domain/adapter"
	"github.com/tidbcloud/tidb-management-service/serverless-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/httputil"
)

const (
	defaultSys9TokenTTLSec       = 900
	defaultSys9SessionRefreshSec = 120
)

type sys9AuthManager struct {
	client         *resty.Client
	secretProvider adapter.Sys9SecretProvider

	mu      sync.RWMutex
	cache   map[string]*adapter.Sys9Session
	sfGroup singleflight.Group
}

func NewSys9AuthManager(secretProvider adapter.Sys9SecretProvider) adapter.Sys9AuthManager {
	return &sys9AuthManager{
		client:         resty.New(),
		secretProvider: secretProvider,
		cache:          make(map[string]*adapter.Sys9Session),
	}
}

func (m *sys9AuthManager) GetSession(
	ctx context.Context,
	region *entity.Region,
) (*adapter.Sys9Session, error) {
	if region == nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "region is required")
	}
	if region.GetServiceAddress() == "" {
		return nil, bizerr.New(bizerr.ErrCodeParams, "service_address is required")
	}
	if region.GetSys9Config() == nil || region.GetSys9Config().GetOemProfileRef() == "" {
		return nil, bizerr.New(bizerr.ErrCodeParams, "sys9 oem profile ref is required")
	}

	key := m.cacheKey(region)
	now := time.Now()
	if session := m.getCachedSession(key); session != nil {
		if now.Before(session.RefreshAfter) {
			return cloneSession(session), nil
		}
		if now.Before(session.ExpiresAt) {
			refreshCtx := context.WithoutCancel(ctx)
			go func() {
				_, _ = m.refreshWithSingleflight(refreshCtx, key, region)
			}()
			return cloneSession(session), nil
		}
	}

	sess, err := m.refreshWithSingleflight(ctx, key, region)
	if err != nil {
		return nil, err
	}
	return cloneSession(sess), nil
}

func (m *sys9AuthManager) Invalidate(_ context.Context, region *entity.Region) {
	if region == nil {
		return
	}
	key := m.cacheKey(region)
	m.mu.Lock()
	delete(m.cache, key)
	m.mu.Unlock()
}

func (m *sys9AuthManager) refreshWithSingleflight(
	ctx context.Context,
	key string,
	region *entity.Region,
) (*adapter.Sys9Session, error) {
	v, err, _ := m.sfGroup.Do(key, func() (any, error) {
		return m.createSession(ctx, region)
	})
	if err != nil {
		return nil, err
	}
	s, ok := v.(*adapter.Sys9Session)
	if !ok {
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "unexpected session type")
	}
	m.mu.Lock()
	m.cache[key] = cloneSession(s)
	m.mu.Unlock()
	return s, nil
}

func (m *sys9AuthManager) createSession(
	ctx context.Context,
	region *entity.Region,
) (*adapter.Sys9Session, error) {
	profile, err := m.secretProvider.ResolveProfile(ctx, region.GetSys9Config().GetOemProfileRef())
	if err != nil {
		return nil, err
	}

	challengeID, err := m.createChallenge(ctx, region, profile)
	if err != nil {
		return nil, err
	}
	if challengeID == "" {
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "empty challenge id")
	}

	sessionToken, accessToken, expiresIn, err := m.exchangeToken(ctx, region.GetServiceAddress(), profile, challengeID)
	if err != nil {
		return nil, err
	}
	if sessionToken == "" && accessToken == "" {
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "empty token response")
	}

	ttlSec := int(region.GetSys9Config().GetTokenTtlSec())
	if ttlSec <= 0 {
		if expiresIn > 0 {
			ttlSec = expiresIn
		} else {
			ttlSec = defaultSys9TokenTTLSec
		}
	}
	refreshSec := int(region.GetSys9Config().GetSessionRefreshBeforeSec())
	if refreshSec <= 0 {
		refreshSec = defaultSys9SessionRefreshSec
	}

	now := time.Now()
	expiresAt := now.Add(time.Duration(ttlSec) * time.Second)
	refreshAfter := expiresAt.Add(-time.Duration(refreshSec) * time.Second)
	if refreshAfter.Before(now) {
		refreshAfter = now
	}

	return &adapter.Sys9Session{
		SessionToken:     sessionToken,
		AccessToken:      accessToken,
		RequestContextID: uuid.NewString(),
		ExpiresAt:        expiresAt,
		RefreshAfter:     refreshAfter,
	}, nil
}

func (m *sys9AuthManager) createChallenge(
	ctx context.Context,
	region *entity.Region,
	profile *adapter.Sys9OEMProfile,
) (string, error) {
	if region == nil {
		return "", bizerr.New(bizerr.ErrCodeParams, "region is required")
	}
	serviceAddress := region.GetServiceAddress()
	if serviceAddress == "" {
		return "", bizerr.New(bizerr.ErrCodeParams, "service_address is required")
	}
	apiPath := "/api/v1/federation/bootstrap-challenges"
	temporaryDPoPPublicKey := profile.TemporaryDPoPPublicKey
	if temporaryDPoPPublicKey == "" {
		temporaryDPoPPublicKey = profile.DPoPPublicKey
	}
	body := map[string]any{
		"brandId":                profile.BrandID,
		"issuer":                 profile.Issuer,
		"temporaryDpopPublicKey": temporaryDPoPPublicKey,
	}
	if len(profile.EdgeBinding) > 0 {
		body["edgeBinding"] = profile.EdgeBinding
	}
	if profile.Audience != "" && strings.ToLower(region.GetSys9Config().GetAuthMode()) != "mock" {
		body["audience"] = profile.Audience
	}
	if profile.TLSExporterDigest != "" {
		body["tlsExporterSha256"] = profile.TLSExporterDigest
	}

	resp, err := httputil.DoRequest(ctx, m.client, resty.MethodPost, serviceAddress, apiPath, func(req *resty.Request) {
		req.ForceContentType("application/json")
		req.SetBody(body)
	})
	if err != nil {
		return "", bizerr.New(bizerr.ErrCodeUnknown, "create sys9 challenge failed: %v", err)
	}
	if resp.StatusCode() >= 400 {
		return "", mapSys9Error(resp.StatusCode(), resp.Body())
	}

	var out struct {
		ChallengeID string `json:"challengeId"`
		Challenge   string `json:"challenge"`
		ID          string `json:"id"`
	}
	if err := json.Unmarshal(resp.Body(), &out); err != nil {
		return "", bizerr.New(bizerr.ErrCodeUnknown, "decode sys9 challenge failed: %v", err)
	}
	if out.ChallengeID != "" {
		return out.ChallengeID, nil
	}
	if out.Challenge != "" {
		return out.Challenge, nil
	}
	return out.ID, nil
}

func (m *sys9AuthManager) exchangeToken(
	ctx context.Context,
	serviceAddress string,
	profile *adapter.Sys9OEMProfile,
	challengeID string,
) (sessionToken, accessToken string, expiresInSec int, err error) {
	apiPath := "/api/v1/tokens:exchange"
	requestContextID := uuid.NewString()
	body := map[string]any{
		"brandId":              profile.BrandID,
		"challengeId":          challengeID,
		"federationAssertion":  profile.Assertion,
		"dpopProof":            profile.DPoPProof,
		"subjectAnchor":        profile.SubjectAnchor,
		"subjectAnchorVersion": profile.SubjectAnchorVersion,
		"entityAnchor":         profile.EntityAnchor,
		"entityAnchorVersion":  profile.EntityAnchorVersion,
	}
	if len(profile.EdgeBinding) > 0 {
		body["edgeBinding"] = profile.EdgeBinding
	}

	resp, err := httputil.DoRequest(ctx, m.client, resty.MethodPost, serviceAddress, apiPath, func(req *resty.Request) {
		req.ForceContentType("application/json")
		req.SetHeader("X-Request-Context-ID", requestContextID)
		req.SetBody(body)
	})
	if err != nil {
		return "", "", 0, bizerr.New(bizerr.ErrCodeUnknown, "exchange sys9 token failed: %v", err)
	}
	if resp.StatusCode() >= 400 {
		return "", "", 0, mapSys9Error(resp.StatusCode(), resp.Body())
	}

	var out struct {
		SessionToken string `json:"sessionToken"`
		AccessToken  string `json:"accessToken"`
		Token        string `json:"token"`
		SessionID    string `json:"sessionId"`
		ExpiresInSec int    `json:"expiresInSec"`
		ExpiresIn    int    `json:"expiresIn"`
	}
	if err := json.Unmarshal(resp.Body(), &out); err != nil {
		return "", "", 0, bizerr.New(bizerr.ErrCodeUnknown, "decode sys9 exchange response failed: %v", err)
	}

	if out.SessionToken != "" {
		sessionToken = out.SessionToken
	} else {
		sessionToken = out.SessionID
	}

	if out.AccessToken != "" {
		accessToken = out.AccessToken
	} else {
		accessToken = out.Token
	}
	if accessToken == "" {
		accessToken = sessionToken
	}
	if out.ExpiresInSec > 0 {
		expiresInSec = out.ExpiresInSec
	} else {
		expiresInSec = out.ExpiresIn
	}

	return sessionToken, accessToken, expiresInSec, nil
}

func (m *sys9AuthManager) cacheKey(region *entity.Region) string {
	authMode := ""
	oemProfileRef := ""
	if region.GetSys9Config() != nil {
		authMode = region.GetSys9Config().GetAuthMode()
		oemProfileRef = region.GetSys9Config().GetOemProfileRef()
	}
	return strings.Join([]string{
		region.GetCloudProviderId().String(),
		region.GetId().String(),
		region.GetServicePlan().String(),
		region.GetEngineType().String(),
		authMode,
		oemProfileRef,
	}, "|")
}

func (m *sys9AuthManager) getCachedSession(key string) *adapter.Sys9Session {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return cloneSession(m.cache[key])
}

func cloneSession(s *adapter.Sys9Session) *adapter.Sys9Session {
	if s == nil {
		return nil
	}
	cloned := *s
	return &cloned
}
