package main

import (
	"github.com/pkg/errors"
	"gopkg.in/square/go-jose.v2"
	"log"
	"net/url"
	"time"
)

const (
	REQUEST_BUFFER_SIZE      = 10
	KEY_MAP_INITIAL_SIZE     = 5
	DEFAULT_REFRESH_INTERVAL = 15 * time.Minute
	MIN_REFRESH_INTERVAL     = 1 * time.Minute
)

type KeyRequest struct {
	KeyId    string
	Response chan *jose.JSONWebKey
}

type JWKSFetcher interface {
	Run()
	Fetch() error
	GetKey(string) (*jose.JSONWebKey, error)
	Done()
}

type jwksFetcher struct {
	keyMap     map[string]jose.JSONWebKey
	httpClient CautiousHTTPClient
	validator  KeyValidator
	fetchTimer *time.Timer
	url        *url.URL
	requests   chan *KeyRequest
	done       chan bool
}

func NewJWKSFetcher(h CautiousHTTPClient, url *url.URL, issuer string, root string) JWKSFetcher {
	val := NewKeyValidator(HostFromURL(issuer))
	val.LoadRootPEM(root)

	return &jwksFetcher{
		httpClient: h,
		validator:  val,
		url:        url,
		fetchTimer: time.NewTimer(DEFAULT_REFRESH_INTERVAL),
		requests:   make(chan *KeyRequest, REQUEST_BUFFER_SIZE),
		keyMap:     make(map[string]jose.JSONWebKey, KEY_MAP_INITIAL_SIZE),
		done:       make(chan bool),
	}
}

func (f *jwksFetcher) Fetch() error {
	var jwks jose.JSONWebKeySet
	timeout, err := f.httpClient.GetJSONExpires(f.url.String(), &jwks)
	if err != nil {
		return errors.WithStack(err)
	}

	for _, k := range jwks.Keys {
		err = f.validator.Validate(k)
		if err == nil {
			f.keyMap[k.KeyID] = k
		} else {
			log.Printf("Rejecting key %q because %q", k.KeyID, err)
		}
	}

	if timeout < MIN_REFRESH_INTERVAL {
		timeout = MIN_REFRESH_INTERVAL
	}

	success := f.fetchTimer.Reset(timeout)
	if !success {
		f.fetchTimer = time.NewTimer(timeout)
	}

	return nil
}

func (f *jwksFetcher) Run() {
	for {
		select {
		// Incoming request for a key, return key or nil in no key
		case r := <-f.requests:
			if v, ok := f.keyMap[r.KeyId]; ok {
				r.Response <- &v
			} else {
				r.Response <- nil
			}
		case <-f.fetchTimer.C:
			f.Fetch()
		case <-f.done:
			return
		}
	}
}

func (f *jwksFetcher) Done() {
	f.done <- true
}

func (f *jwksFetcher) GetKey(kid string) (*jose.JSONWebKey, error) {
	r := &KeyRequest{
		KeyId:    kid,
		Response: make(chan *jose.JSONWebKey),
	}

	f.requests <- r

	if res := <-r.Response; res == nil {
		return nil, errors.Errorf("Key not found for ID")
	} else {
		return res, nil
	}
}