package main

import (
	"encoding/json"
	"log"
	"net/http"
	"net/url"
)

type tokenRequest struct {
	TokenEndpoint string `json:"token_endpoint"`
	Issuer        string `json:"issuer"`
	KeyID         string `json:"key_id,omitempty"`
	GrantType     string `json:"grant_type"`
	Code          string `json:"code,omitempty"`
	RedirectURI   string `json:"redirect_uri,omitempty"`
	CodeVerifier  string `json:"code_verifier,omitempty"`
	RefreshToken  string `json:"refresh_token,omitempty"`
}

func HandleToken(signers *SignerSet, clientID string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var req tokenRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeJSONError(w, http.StatusBadRequest, "invalid_request", "invalid JSON body")
			return
		}

		if req.TokenEndpoint == "" {
			writeJSONError(w, http.StatusBadRequest, "invalid_request", "token_endpoint is required")
			return
		}
		if req.Issuer == "" {
			writeJSONError(w, http.StatusBadRequest, "invalid_request", "issuer is required")
			return
		}
		if req.GrantType == "" {
			writeJSONError(w, http.StatusBadRequest, "invalid_request", "grant_type is required")
			return
		}

		if err := ValidateTokenEndpointForIssuer(r.Context(), req.Issuer, req.TokenEndpoint); err != nil {
			writeAPIError(w, err)
			return
		}

		candidateKeyIDs, err := signers.CandidateKeyIDs(req.KeyID)
		if err != nil {
			writeJSONError(w, http.StatusBadRequest, "invalid_request", err.Error())
			return
		}

		params := url.Values{}
		params.Set("grant_type", req.GrantType)
		params.Set("client_id", clientID)

		if req.Code != "" {
			params.Set("code", req.Code)
		}
		if req.RedirectURI != "" {
			params.Set("redirect_uri", req.RedirectURI)
		}
		if req.CodeVerifier != "" {
			params.Set("code_verifier", req.CodeVerifier)
		}
		if req.RefreshToken != "" {
			params.Set("refresh_token", req.RefreshToken)
		}

		dpopHeader := r.Header.Get("DPoP")

		var proxied *upstreamResponse
		var usedKeyID string

		for i, keyID := range candidateKeyIDs {
			signer, err := signers.Lookup(keyID)
			if err != nil {
				writeJSONError(w, http.StatusBadRequest, "invalid_request", err.Error())
				return
			}

			assertion, err := GenerateClientAssertion(signer, clientID, req.Issuer)
			if err != nil {
				log.Printf("failed to generate client assertion: %v", err)
				writeJSONError(w, http.StatusInternalServerError, "server_error", "failed to generate client assertion")
				return
			}

			attemptParams := cloneValues(params)
			attemptParams.Set("client_assertion_type", "urn:ietf:params:oauth:client-assertion-type:jwt-bearer")
			attemptParams.Set("client_assertion", assertion)

			proxied, err = PostForm(r.Context(), req.TokenEndpoint, attemptParams, dpopHeader)
			if err != nil {
				log.Printf("proxy request failed: %v", err)
				writeAPIError(w, upstreamRequestError("upstream request failed"))
				return
			}

			usedKeyID = keyID
			if req.KeyID == "" && i < len(candidateKeyIDs)-1 && isInvalidClientResponse(proxied) {
				continue
			}

			break
		}

		w.Header().Set(authProxyKeyIDHeader, usedKeyID)
		if err := WriteProxiedResponse(w, proxied); err != nil {
			log.Printf("failed to write proxied response: %v", err)
		}
	}
}

func isInvalidClientResponse(resp *upstreamResponse) bool {
	if resp == nil {
		return false
	}
	if resp.statusCode != http.StatusBadRequest && resp.statusCode != http.StatusUnauthorized {
		return false
	}

	var payload struct {
		Error string `json:"error"`
	}
	if err := json.Unmarshal(resp.body, &payload); err != nil {
		return false
	}

	return payload.Error == "invalid_client"
}

func cloneValues(values url.Values) url.Values {
	cloned := make(url.Values, len(values))
	for key, entries := range values {
		cloned[key] = append([]string(nil), entries...)
	}
	return cloned
}