package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"slices"
	"sync"
	"time"
)

const (
	authServerMetadataCacheTTL  = 5 * time.Minute
	maxMetadataResponseBytes    = 256 << 10
	maxUpstreamResponseBodySize = 1 << 20
)

type authServerMetadata struct {
	Issuer                                string   `json:"issuer"`
	TokenEndpoint                         string   `json:"token_endpoint"`
	TokenEndpointAuthMethodsSupported     []string `json:"token_endpoint_auth_methods_supported"`
	TokenEndpointAuthSigningAlgsSupported []string `json:"token_endpoint_auth_signing_alg_values_supported"`
	RequirePushedAuthorizationRequests    bool     `json:"require_pushed_authorization_requests"`
	PushedAuthorizationRequestEndpoint    string   `json:"pushed_authorization_request_endpoint"`
}

type cachedAuthServerMetadata struct {
	metadata  *authServerMetadata
	expiresAt time.Time
}

var (
	authServerMetadataCache = struct {
		mu      sync.Mutex
		entries map[string]cachedAuthServerMetadata
	}{
		entries: make(map[string]cachedAuthServerMetadata),
	}

	metadataClient = newPublicHTTPClient(10 * time.Second)
)

func clearAuthServerMetadataCache() {
	authServerMetadataCache.mu.Lock()
	defer authServerMetadataCache.mu.Unlock()
	authServerMetadataCache.entries = make(map[string]cachedAuthServerMetadata)
}

func ResolveAuthServerMetadata(ctx context.Context, issuer string) (*authServerMetadata, error) {
	issuerURL, err := ValidateIssuer(issuer)
	if err != nil {
		return nil, invalidRequestError("invalid issuer")
	}

	now := time.Now()

	authServerMetadataCache.mu.Lock()
	entry, ok := authServerMetadataCache.entries[issuer]
	if ok && now.Before(entry.expiresAt) {
		authServerMetadataCache.mu.Unlock()
		return entry.metadata, nil
	}
	authServerMetadataCache.mu.Unlock()

	metadataURL := issuerURL.ResolveReference(&url.URL{Path: "/.well-known/oauth-authorization-server"})
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, metadataURL.String(), nil)
	if err != nil {
		return nil, upstreamRequestError("failed to create metadata request")
	}

	resp, err := metadataClient.Do(req)
	if err != nil {
		return nil, upstreamRequestError("failed to fetch authorization server metadata")
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, upstreamRequestError(fmt.Sprintf("authorization server metadata returned HTTP %d", resp.StatusCode))
	}

	body, err := io.ReadAll(io.LimitReader(resp.Body, maxMetadataResponseBytes+1))
	if err != nil {
		return nil, upstreamRequestError("failed to read authorization server metadata")
	}
	if len(body) > maxMetadataResponseBytes {
		return nil, upstreamRequestError("authorization server metadata response was too large")
	}

	var metadata authServerMetadata
	if err := json.Unmarshal(body, &metadata); err != nil {
		return nil, upstreamRequestError("authorization server metadata was not valid JSON")
	}

	if err := metadata.Validate(issuer); err != nil {
		return nil, invalidRequestError(err.Error())
	}

	authServerMetadataCache.mu.Lock()
	authServerMetadataCache.entries[issuer] = cachedAuthServerMetadata{
		metadata:  &metadata,
		expiresAt: now.Add(authServerMetadataCacheTTL),
	}
	authServerMetadataCache.mu.Unlock()

	return &metadata, nil
}

func (m *authServerMetadata) Validate(expectedIssuer string) error {
	if _, err := ValidateIssuer(m.Issuer); err != nil {
		return fmt.Errorf("issuer metadata contained an invalid issuer")
	}
	if m.Issuer != expectedIssuer {
		return fmt.Errorf("issuer metadata did not match the requested issuer")
	}
	if _, err := validateEndpointURL(m.TokenEndpoint); err != nil {
		return fmt.Errorf("issuer metadata contained an invalid token_endpoint")
	}
	if !slices.Contains(m.TokenEndpointAuthMethodsSupported, "private_key_jwt") {
		return fmt.Errorf("issuer metadata does not support private_key_jwt")
	}
	if !slices.Contains(m.TokenEndpointAuthSigningAlgsSupported, "ES256") {
		return fmt.Errorf("issuer metadata does not support ES256 client assertions")
	}
	if m.PushedAuthorizationRequestEndpoint != "" {
		if _, err := validateEndpointURL(m.PushedAuthorizationRequestEndpoint); err != nil {
			return fmt.Errorf("issuer metadata contained an invalid pushed_authorization_request_endpoint")
		}
	}
	if m.RequirePushedAuthorizationRequests && m.PushedAuthorizationRequestEndpoint == "" {
		return fmt.Errorf("issuer metadata requires PAR but does not advertise a PAR endpoint")
	}

	return nil
}

func ValidateTokenEndpointForIssuer(ctx context.Context, issuer string, tokenEndpoint string) error {
	if err := ValidateTokenEndpoint(tokenEndpoint); err != nil {
		return invalidRequestError("invalid token_endpoint")
	}

	metadata, err := ResolveAuthServerMetadata(ctx, issuer)
	if err != nil {
		return err
	}

	if metadata.TokenEndpoint != tokenEndpoint {
		return invalidRequestError("token_endpoint does not match issuer metadata")
	}

	return nil
}

func ValidatePAREndpointForIssuer(ctx context.Context, issuer string, parEndpoint string) error {
	if err := ValidatePAREndpoint(parEndpoint); err != nil {
		return invalidRequestError("invalid par_endpoint")
	}

	metadata, err := ResolveAuthServerMetadata(ctx, issuer)
	if err != nil {
		return err
	}

	if metadata.PushedAuthorizationRequestEndpoint == "" {
		return invalidRequestError("issuer does not advertise a pushed_authorization_request_endpoint")
	}
	if metadata.PushedAuthorizationRequestEndpoint != parEndpoint {
		return invalidRequestError("par_endpoint does not match issuer metadata")
	}

	return nil
}