package main

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/envoyproxy/envoy/contrib/golang/common/go/api"
	"github.com/golang-jwt/jwt/v5"
)

type filter struct {
	api.PassThroughStreamFilter
	callbacks api.FilterCallbackHandler
	config    *Config
}

// *filter.DecodeHeaders, *filter.DecodeData, *filter.DecodeTrailers, *filter.EncodeData and *filter.EncodeTrailers
// are omitted for brevity, they all just return api.Continue.

// *filter.OnLog, *filter.OnLogDownstreamStart, *filter.OnLogDownstreamPeriodic and *filter.OnDestroy
// are omitted for brevity, they are all just empty.

func (f *filter) EncodeHeaders(header api.ResponseHeaderMap, endStream bool) api.StatusType {
	setCookies := header.Values("set-cookie")
	if len(setCookies) == 0 {
		return api.Continue
	}

	header.Del("set-cookie")
	for key, value := range setCookies {
		if !strings.HasPrefix(value, "IdToken=") {
			continue
		}
		parts := strings.Split(value, ";")
		_, cookieValue, _ := strings.Cut(parts[0], "=")
		expiry, err := readIDTokenExpiry(cookieValue)
		if err != nil {
			api.LogErrorf("read IdToken expiry: %s", err)
			break
		}
		newMaxAge := int(time.Until(expiry.Time).Seconds())
		api.LogDebugf("fixing max-age - new lifetime: %d", newMaxAge)
		setCookies[key] = replaceMaxAgeInSetCookieValue(value, newMaxAge)
		break
	}
	for _, v := range setCookies {
		header.Add("set-cookie", v)
	}

	return api.Continue
}

func readIDTokenExpiry(idToken string) (*jwt.NumericDate, error) {
	unverifiedToken, _, err := new(jwt.Parser).ParseUnverified(idToken, jwt.MapClaims{})
	if err != nil {
		return nil, fmt.Errorf("parse the JWT: %s", err)
	}
	if unverifiedClaims, ok := unverifiedToken.Claims.(jwt.MapClaims); ok {
		expiry, err := unverifiedClaims.GetExpirationTime()
		if err != nil {
			return nil, fmt.Errorf("get expiry: %s", err)
		}
		return expiry, nil
	}
	return nil, errors.New("invalid JWT token")
}

func replaceMaxAgeInSetCookieValue(value string, maxAge int) string {
	return regexp.MustCompile(`Max-Age=([0-9]+)`).ReplaceAllString(value, fmt.Sprintf("Max-Age=%d", maxAge))
}