package token

import (
	"context"
	"fmt"

	"github.com/tidbcloud/aisupport/packages/sdk-go/infra/utils/container/gmap"
	"golang.org/x/oauth2"
)

type TokenStore interface {
	Save(ctx context.Context, user_email string, token *oauth2.Token) error
	Get(ctx context.Context, user_email string) (*oauth2.Token, error)
}

var _ TokenStore = (*mem)(nil)

type mem struct {
	tokens *gmap.Map[string, *oauth2.Token]
}

func NewMem() TokenStore {
	return &mem{
		tokens: gmap.NewMap[string, *oauth2.Token](),
	}
}

func (s *mem) Save(ctx context.Context, user_email string, token *oauth2.Token) error {
	if user_email == "" {
		return fmt.Errorf("user email is empty")
	}
	if token == nil {
		return fmt.Errorf("token is nil")
	}

	s.tokens.Store(user_email, cloneToken(token))
	return nil
}

func (s *mem) Get(ctx context.Context, user_email string) (*oauth2.Token, error) {
	if user_email == "" {
		return nil, fmt.Errorf("user email is empty")
	}

	token, ok := s.tokens.Load(user_email)
	if !ok {
		return nil, nil
	}
	return cloneToken(token), nil
}

func cloneToken(token *oauth2.Token) *oauth2.Token {
	if token == nil {
		return nil
	}
	cloned := *token
	return &cloned
}