package workloadapi

import (
	"context"
	"sync"

	"github.com/spiffe/go-spiffe/v2/bundle/jwtbundle"
	"github.com/spiffe/go-spiffe/v2/spiffeid"
	"github.com/spiffe/go-spiffe/v2/svid/jwtsvid"
	"github.com/zeebo/errs"
)

var jwtsourceErr = errs.Class("jwtsource")

// JWTSource is a source of JWT-SVID and JWT bundles maintained via the
// Workload API.
type JWTSource struct {
	watcher *watcher

	mtx     sync.RWMutex
	bundles *jwtbundle.Set

	closeMtx sync.RWMutex
	closed   bool
}

// NewJWTSource creates a new JWTSource. It blocks until the initial update
// has been received from the Workload API.
func NewJWTSource(ctx context.Context, options ...JWTSourceOption) (_ *JWTSource, err error) {
	config := &jwtSourceConfig{}
	for _, option := range options {
		option.configureJWTSource(config)
	}

	s := &JWTSource{}

	s.watcher, err = newWatcher(ctx, config.watcher, nil, s.setJWTBundles)
	if err != nil {
		return nil, err
	}

	return s, nil
}

// Close closes the source, dropping the connection to the Workload API.
// Other source methods will return an error after Close has been called.
// The underlying Workload API client will also be closed if it is owned by
// the JWTSource (i.e. not provided via the WithClient option).
func (s *JWTSource) Close() error {
	s.closeMtx.Lock()
	s.closed = true
	s.closeMtx.Unlock()

	return s.watcher.Close()
}

// FetchJWTSVID fetches a JWT-SVID from the source with the given parameters.
// It implements the jwtsvid.Source interface.
func (s *JWTSource) FetchJWTSVID(ctx context.Context, params jwtsvid.Params) (*jwtsvid.SVID, error) {
	if err := s.checkClosed(); err != nil {
		return nil, err
	}
	return s.watcher.client.FetchJWTSVID(ctx, params)
}

// GetJWTBundleForTrustDomain returns the JWT bundle for the given trust
// domain. It implements the jwtbundle.Source interface.
func (s *JWTSource) GetJWTBundleForTrustDomain(trustDomain spiffeid.TrustDomain) (*jwtbundle.Bundle, error) {
	if err := s.checkClosed(); err != nil {
		return nil, err
	}
	return s.bundles.GetJWTBundleForTrustDomain(trustDomain)
}

// WaitUntilUpdated waits until the source is updated or the context is done,
// in which case ctx.Err() is returned.
func (s *JWTSource) WaitUntilUpdated(ctx context.Context) error {
	return s.watcher.WaitUntilUpdated(ctx)
}

// Updated returns a channel that is sent on whenever the source is updated.
func (s *JWTSource) Updated() <-chan struct{} {
	return s.watcher.Updated()
}

func (s *JWTSource) setJWTBundles(bundles *jwtbundle.Set) {
	s.mtx.Lock()
	defer s.mtx.Unlock()
	s.bundles = bundles
}

func (s *JWTSource) checkClosed() error {
	s.closeMtx.RLock()
	defer s.closeMtx.RUnlock()
	if s.closed {
		return jwtsourceErr.New("source is closed")
	}
	return nil
}
