package authapi

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"math/big"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

// signingKey is generated once per test binary; RSA generation is expensive.
var signingKey = sync.OnceValue(func() *rsa.PrivateKey {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(err)
	}
	return key
})

// fakeIdP is a minimal OIDC provider: discovery document, JWKS and token endpoint.
// Tests drive its behaviour through the exported fields before calling the callback.
type fakeIdP struct {
	server   *httptest.Server
	key      *rsa.PrivateKey
	clientID string

	mu                sync.Mutex
	expectedChallenge string
	nonce             string
	subject           string
	groups            []string
	tokenFails        bool
	omitIDToken       bool
	issuerOverride    string
	verifierSeen      string
}

func newFakeIdP(t *testing.T, clientID string) *fakeIdP {
	t.Helper()
	idp := &fakeIdP{key: signingKey(), clientID: clientID, subject: "user-1", groups: []string{"pulse-operator"}}
	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", idp.discovery)
	mux.HandleFunc("/jwks", idp.jwks)
	mux.HandleFunc("/token", idp.token)
	mux.HandleFunc("/authorize", func(response http.ResponseWriter, _ *http.Request) {
		response.WriteHeader(http.StatusOK)
	})
	idp.server = httptest.NewServer(mux)
	t.Cleanup(idp.server.Close)
	return idp
}

func (idp *fakeIdP) configure(mutate func(*fakeIdP)) {
	idp.mu.Lock()
	defer idp.mu.Unlock()
	mutate(idp)
}

func (idp *fakeIdP) codeVerifier() string {
	idp.mu.Lock()
	defer idp.mu.Unlock()
	return idp.verifierSeen
}

func (idp *fakeIdP) discovery(response http.ResponseWriter, _ *http.Request) {
	writeJSON(response, http.StatusOK, map[string]any{
		"issuer":                                idp.server.URL,
		"authorization_endpoint":                idp.server.URL + "/authorize",
		"token_endpoint":                        idp.server.URL + "/token",
		"jwks_uri":                              idp.server.URL + "/jwks",
		"response_types_supported":              []string{"code"},
		"subject_types_supported":               []string{"public"},
		"id_token_signing_alg_values_supported": []string{"RS256"},
	})
}

func (idp *fakeIdP) jwks(response http.ResponseWriter, _ *http.Request) {
	public := &idp.key.PublicKey
	writeJSON(response, http.StatusOK, map[string]any{"keys": []map[string]any{{
		"kty": "RSA",
		"kid": "test-key",
		"alg": "RS256",
		"use": "sig",
		"n":   base64.RawURLEncoding.EncodeToString(public.N.Bytes()),
		"e":   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(public.E)).Bytes()),
	}}})
}

func (idp *fakeIdP) token(response http.ResponseWriter, request *http.Request) {
	if err := request.ParseForm(); err != nil {
		writeJSON(response, http.StatusBadRequest, map[string]any{"error": "invalid_request"})
		return
	}
	idp.mu.Lock()
	defer idp.mu.Unlock()
	idp.verifierSeen = request.PostForm.Get("code_verifier")
	if idp.tokenFails {
		writeJSON(response, http.StatusBadRequest, map[string]any{"error": "invalid_grant"})
		return
	}
	if idp.expectedChallenge != "" {
		digest := sha256.Sum256([]byte(idp.verifierSeen))
		if base64.RawURLEncoding.EncodeToString(digest[:]) != idp.expectedChallenge {
			writeJSON(response, http.StatusBadRequest, map[string]any{"error": "invalid_grant"})
			return
		}
	}
	body := map[string]any{"access_token": "opaque-access-token", "token_type": "Bearer", "expires_in": 3600}
	if !idp.omitIDToken {
		issuer := idp.server.URL
		if idp.issuerOverride != "" {
			issuer = idp.issuerOverride
		}
		now := time.Now()
		body["id_token"] = idp.sign(map[string]any{
			"iss":    issuer,
			"aud":    idp.clientID,
			"sub":    idp.subject,
			"iat":    now.Unix(),
			"exp":    now.Add(5 * time.Minute).Unix(),
			"nonce":  idp.nonce,
			"groups": idp.groups,
		})
	}
	writeJSON(response, http.StatusOK, body)
}

func (idp *fakeIdP) sign(claims map[string]any) string {
	segments := []string{encodeSegment(map[string]any{"alg": "RS256", "typ": "JWT", "kid": "test-key"}), encodeSegment(claims)}
	input := strings.Join(segments, ".")
	digest := sha256.Sum256([]byte(input))
	signature, err := rsa.SignPKCS1v15(rand.Reader, idp.key, crypto.SHA256, digest[:])
	if err != nil {
		panic(err)
	}
	return input + "." + base64.RawURLEncoding.EncodeToString(signature)
}

func encodeSegment(value map[string]any) string {
	encoded, err := json.Marshal(value)
	if err != nil {
		panic(err)
	}
	return base64.RawURLEncoding.EncodeToString(encoded)
}

func writeJSON(response http.ResponseWriter, status int, body any) {
	response.Header().Set("Content-Type", "application/json")
	response.WriteHeader(status)
	_ = json.NewEncoder(response).Encode(body)
}