package oidc_test

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

	"keycape/internal/domain"
	"keycape/internal/server/oidc"
	"keycape/internal/server/telemetry"
)

func verifyWithJWKS(t *testing.T, h *oidc.TokenHandler, token string) {
	t.Helper()
	keys := oidc.NewKeySet()
	keys.AddKey("key-1", &h.SigningKey.PublicKey)
	w := httptest.NewRecorder()
	oidc.NewJWKSHandler(keys).ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/jwks", nil))
	var response struct {
		Keys []oidc.JWK `json:"keys"`
	}
	if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
		t.Fatal(err)
	}
	if len(response.Keys) != 1 {
		t.Fatal("missing signing key")
	}
	key := response.Keys[0]
	n, err := base64.RawURLEncoding.DecodeString(key.N)
	if err != nil {
		t.Fatal(err)
	}
	e, err := base64.RawURLEncoding.DecodeString(key.E)
	if err != nil {
		t.Fatal(err)
	}
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		t.Fatal("invalid JWT")
	}
	sig, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		t.Fatal(err)
	}
	digest := sha256.Sum256([]byte(parts[0] + "." + parts[1]))
	pub := rsa.PublicKey{N: new(big.Int).SetBytes(n), E: int(new(big.Int).SetBytes(e).Int64())}
	if err := rsa.VerifyPKCS1v15(&pub, crypto.SHA256, digest[:], sig); err != nil {
		t.Fatal(err)
	}
}

func TestServiceAudienceIsStaticAndDefaultsToClientID(t *testing.T) {
	for _, audience := range []string{"", "approval-engine"} {
		t.Run(audience, func(t *testing.T) {
			h := serviceTokenHandler(t)
			h.ClientConfig["rapp-qonto"].Audience = audience
			req := tokenRequest(url.Values{"grant_type": {"client_credentials"}, "audience": {"attacker"}, "resource": {"attacker"}})
			req.SetBasicAuth("rapp-qonto", "test-service-secret")
			w := httptest.NewRecorder()
			h.ServeHTTP(w, req)
			if w.Code != 200 {
				t.Fatalf("status %d", w.Code)
			}
			token := decodeTokenResponse(t, w.Body.String())["access_token"].(string)
			want := audience
			if want == "" {
				want = "rapp-qonto"
			}
			if parseJWTPayload(t, token)["aud"] != want {
				t.Fatal("wrong audience")
			}
			verifyWithJWKS(t, h, token)
		})
	}
}

func TestHumanResourceAudiencePreservesIDTokenAudience(t *testing.T) {
	sessions := oidc.NewSessionStore()
	h, _ := newTokenHandler(t, sessions, &mockUserRepo{users: map[string]*domain.User{"alice": aliceUser()}})
	h.ClientConfig["test-client"].Audience = "approval-engine"
	h.ClientConfig["test-client"].AllowedScopes = []string{"openid", "approval:approve"}
	verifier := "test-verifier"
	code := seededSession(sessions, verifier)
	sess, _ := sessions.Get(code)
	sess.Scopes = []string{"openid", "approval:approve"}
	w := httptest.NewRecorder()
	h.ServeHTTP(w, tokenRequest(url.Values{"grant_type": {"authorization_code"}, "client_id": {"test-client"}, "code": {code}, "code_verifier": {verifier}}))
	if w.Code != 200 {
		t.Fatalf("status %d: %s", w.Code, w.Body.String())
	}
	response := decodeTokenResponse(t, w.Body.String())
	access := response["access_token"].(string)
	id := response["id_token"].(string)
	claims := parseJWTPayload(t, access)
	if claims["aud"] != "approval-engine" || claims["scope"] != "openid approval:approve" || claims["principal_type"] != "human" {
		t.Fatalf("wrong access claims: %v", claims)
	}
	if parseJWTPayload(t, id)["aud"] != "test-client" {
		t.Fatal("ID token audience changed")
	}
	verifyWithJWKS(t, h, access)
	verifyWithJWKS(t, h, id)
}

func TestHumanExcessScopeRejectedBeforeAuthentication(t *testing.T) {
	h := newAuthorizeHandler(nil, nil, telemetry.NoopEmitter{})
	params := validAuthorizeParams()
	params.Set("scope", "openid approval:consume")
	w := httptest.NewRecorder()
	h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/authorize?"+params.Encode(), nil))
	if w.Code != http.StatusBadRequest {
		t.Fatalf("status %d", w.Code)
	}
}

func TestTokenRejectsScopeRemovedAfterAuthorization(t *testing.T) {
	sessions := oidc.NewSessionStore()
	h, _ := newTokenHandler(t, sessions, &mockUserRepo{})
	code := seededSession(sessions, "verifier")
	h.ClientConfig["test-client"].AllowedScopes = []string{"openid"}
	w := httptest.NewRecorder()
	h.ServeHTTP(w, tokenRequest(url.Values{"grant_type": {"authorization_code"}, "client_id": {"test-client"}, "code": {code}, "code_verifier": {"verifier"}}))
	if w.Code != http.StatusBadRequest {
		t.Fatalf("status %d", w.Code)
	}
}