package main

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

const forgeInvocationIssuer = "forge/invocation-token"

type forgeValidator struct {
	allowedAppIDs                  map[string]struct{}
	allowedInstallationIDs         map[string]struct{}
	allowUnregisteredInstallations bool
	requirePrincipal               bool
	jwks                           *jwksCache
}

func newForgeValidator(config *Config) (*forgeValidator, error) {
	allowedAppIDs := stringSet(config.AllowedAppIDs)
	if len(allowedAppIDs) == 0 {
		return nil, fmt.Errorf("allowed_app_ids is required")
	}
	allowedInstallationIDs := stringSet(config.AllowedInstallationIDs)
	if len(allowedInstallationIDs) == 0 && !config.AllowUnregisteredInstallations {
		return nil, fmt.Errorf("allowed_installation_ids is required")
	}

	jwksURL := strings.TrimSpace(config.JWKSURL)
	if jwksURL == "" {
		jwksURL = defaultJWKSURL
	}

	requestTimeout := defaultRequestTimeout
	if config.RequestTimeoutMS > 0 {
		requestTimeout = time.Duration(config.RequestTimeoutMS) * time.Millisecond
	}

	return &forgeValidator{
		allowedAppIDs:                  allowedAppIDs,
		allowedInstallationIDs:         allowedInstallationIDs,
		allowUnregisteredInstallations: config.AllowUnregisteredInstallations,
		requirePrincipal:               config.RequirePrincipal,
		jwks: newJWKSCache(jwksURL, &http.Client{
			Timeout: requestTimeout,
		}, defaultJWKSCacheTTL),
	}, nil
}

func (v *forgeValidator) validate(ctx context.Context, rawToken string) (forgeIdentity, error) {
	claims := &forgeInvocationClaims{}
	parsed, err := jwt.ParseWithClaims(
		rawToken,
		claims,
		func(token *jwt.Token) (any, error) {
			if token.Method.Alg() != jwt.SigningMethodRS256.Alg() {
				return nil, fmt.Errorf("unexpected signing algorithm %q", token.Method.Alg())
			}
			keyID, _ := token.Header["kid"].(string)
			return v.jwks.key(ctx, keyID)
		},
		jwt.WithValidMethods([]string{jwt.SigningMethodRS256.Alg()}),
		jwt.WithIssuer(forgeInvocationIssuer),
		jwt.WithExpirationRequired(),
		jwt.WithLeeway(5*time.Second),
	)
	if err != nil {
		return forgeIdentity{}, fmt.Errorf("validate FIT: %w", err)
	}
	if !parsed.Valid {
		return forgeIdentity{}, fmt.Errorf("FIT is invalid")
	}

	appID := strings.TrimSpace(claims.App.ID)
	if _, ok := v.allowedAppIDs[appID]; !ok {
		return forgeIdentity{}, fmt.Errorf("Forge app is not allowed")
	}
	if !containsAudience(claims.Audience, appID) {
		return forgeIdentity{}, fmt.Errorf("FIT audience does not match Forge app")
	}

	installationID := strings.TrimSpace(claims.App.InstallationID)
	_, installationAllowlisted := v.allowedInstallationIDs[installationID]
	if !installationAllowlisted && !v.allowUnregisteredInstallations {
		return forgeIdentity{}, fmt.Errorf("Forge installation is not allowed")
	}

	principal := strings.TrimSpace(claims.Principal)
	if v.requirePrincipal && principal == "" {
		return forgeIdentity{}, fmt.Errorf("FIT principal is required")
	}

	return forgeIdentity{
		AppID:                   appID,
		InstallationAllowlisted: installationAllowlisted,
		InstallationID:          installationID,
		InvocationID:            strings.TrimSpace(claims.ID),
		Principal:               principal,
	}, nil
}

func containsAudience(audience jwt.ClaimStrings, expected string) bool {
	for _, value := range audience {
		if value == expected {
			return true
		}
	}
	return false
}

func stringSet(values []string) map[string]struct{} {
	result := make(map[string]struct{}, len(values))
	for _, value := range values {
		if normalized := strings.TrimSpace(value); normalized != "" {
			result[normalized] = struct{}{}
		}
	}
	return result
}