package api

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

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

func TestOIDCAccessTokenVerificationAndRoleGate(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	var issuer string
	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]string{"issuer": issuer, "jwks_uri": issuer + "/keys"})
	})
	mux.HandleFunc("/keys", func(w http.ResponseWriter, r *http.Request) {
		e := big.NewInt(int64(key.PublicKey.E)).Bytes()
		_ = json.NewEncoder(w).Encode(map[string]any{"keys": []any{map[string]string{
			"kty": "RSA", "use": "sig", "alg": "RS256", "kid": "test-key",
			"n": base64.RawURLEncoding.EncodeToString(key.PublicKey.N.Bytes()),
			"e": base64.RawURLEncoding.EncodeToString(e),
		}}})
	})
	server := httptest.NewServer(mux)
	defer server.Close()
	issuer = server.URL
	a := newAuthenticator(config.Auth{Mode: "oidc", OIDCIssuer: issuer, OIDCAudience: "fancy-dns-api"})
	for _, tc := range []struct {
		name         string
		roles        []string
		wantReadOnly bool
		bad          bool
		badSignature bool
		badIssuer    bool
	}{
		{name: "admin can write", roles: []string{"admin"}, wantReadOnly: false},
		{name: "ordinary account is read only", roles: []string{"staff"}, wantReadOnly: true},
		{name: "wrong audience refused", roles: []string{"admin"}, bad: true},
		{name: "wrong issuer refused", roles: []string{"admin"}, bad: true, badIssuer: true},
		{name: "expired token refused", roles: []string{"admin"}, bad: true},
		{name: "id token refused", roles: []string{"admin"}, bad: true},
		{name: "bad signature refused", roles: []string{"admin"}, bad: true, badSignature: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			aud := "fancy-dns-api"
			use := "access"
			exp := time.Now().Add(time.Hour).Unix()
			if strings.Contains(tc.name, "wrong audience") {
				aud = "another-api"
			}
			if strings.Contains(tc.name, "expired") {
				exp = time.Now().Add(-time.Hour).Unix()
			}
			if strings.Contains(tc.name, "id token") {
				use = "id"
			}
			tokenIssuer := issuer
			if tc.badIssuer {
				tokenIssuer += "/wrong"
			}
			token := signedOIDCToken(t, key, map[string]any{"iss": tokenIssuer, "sub": "user-1", "aud": aud, "exp": exp, "iat": time.Now().Unix(), "token_use": use, "name": "Test User", "roles": tc.roles})
			if tc.badSignature {
				parts := strings.Split(token, ".")
				parts[2] = strings.Repeat("A", len(parts[2]))
				token = strings.Join(parts, ".")
			}
			p, err := a.oidcPrincipal(token)
			if tc.bad {
				if err == nil {
					t.Fatalf("accepted invalid token as %+v", p)
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			if p.Mode != "oidc" || p.Name != "Test User" || p.ReadOnly != tc.wantReadOnly {
				t.Fatalf("principal = %+v", p)
			}
		})
	}
}

func signedOIDCToken(t *testing.T, key *rsa.PrivateKey, claims map[string]any) string {
	t.Helper()
	enc := base64.RawURLEncoding.EncodeToString
	header, err := json.Marshal(map[string]string{"alg": "RS256", "kid": "test-key", "typ": "JWT"})
	if err != nil {
		t.Fatal(err)
	}
	payload, err := json.Marshal(claims)
	if err != nil {
		t.Fatal(err)
	}
	message := enc(header) + "." + enc(payload)
	digest := sha256.Sum256([]byte(message))
	signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
	if err != nil {
		t.Fatal(err)
	}
	return message + "." + enc(signature)
}