package auth

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"
	"time"
)

// --- OIDC ---

// fakeIssuer hosts a discovery document and a JWKS for a generated RSA key,
// and mints RS256 tokens signed by it.
type fakeIssuer struct {
	srv *httptest.Server
	key *rsa.PrivateKey
}

func newFakeIssuer(t *testing.T) *fakeIssuer {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	f := &fakeIssuer{key: key}
	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]string{"jwks_uri": f.srv.URL + "/jwks"})
	})
	mux.HandleFunc("/jwks", func(w http.ResponseWriter, _ *http.Request) {
		pub := &key.PublicKey
		_ = json.NewEncoder(w).Encode(map[string]any{
			"keys": []map[string]string{{
				"kty": "RSA",
				"kid": "test-key",
				"n":   base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
				"e":   base64.RawURLEncoding.EncodeToString(mustBigIntBytes(pub.E)),
			}},
		})
	})
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	return f
}

// mint signs an RS256 JWT with the given claims (kid "test-key").
func (f *fakeIssuer) mint(t *testing.T, claims map[string]any) string {
	t.Helper()
	enc := func(v any) string {
		b, err := json.Marshal(v)
		if err != nil {
			t.Fatal(err)
		}
		return base64.RawURLEncoding.EncodeToString(b)
	}
	signing := enc(map[string]string{"alg": "RS256", "kid": "test-key"}) + "." + enc(claims)
	digest := sha256.Sum256([]byte(signing))
	sig, err := rsa.SignPKCS1v15(rand.Reader, f.key, crypto.SHA256, digest[:])
	if err != nil {
		t.Fatal(err)
	}
	return signing + "." + base64.RawURLEncoding.EncodeToString(sig)
}

// mustBigIntBytes turns an RSA exponent (typically 65537) into its
// big-endian byte encoding for the JWK "e" field.
func mustBigIntBytes(e int) []byte {
	// smallest big-endian encoding of e.
	n := e
	if n == 0 {
		return []byte{0}
	}
	var buf []byte
	for n > 0 {
		buf = append([]byte{byte(n & 0xff)}, buf...)
		n >>= 8
	}
	return buf
}

func bearerReq(token string) *http.Request {
	r := httptest.NewRequest("GET", "/api/nodes", nil)
	r.Header.Set("Authorization", "Bearer "+token)
	return r
}

func TestOIDCAuthenticate(t *testing.T) {
	iss := newFakeIssuer(t)
	o := NewOIDC(iss.srv.URL, "my-client")

	valid := map[string]any{
		"iss":   iss.srv.URL,
		"sub":   "user-1",
		"email": "user-1@example.com",
		"name":  "User One",
		"aud":   "my-client",
		"exp":   time.Now().Add(time.Hour).Unix(),
		"scope": "openid profile",
	}

	p, err := o.Authenticate(bearerReq(iss.mint(t, valid)))
	if err != nil {
		t.Fatalf("valid token rejected: %v", err)
	}
	if p.Subject != "user-1" || p.Kind != KindOIDC || !p.HasScope("profile") {
		t.Fatalf("unexpected principal: %+v", p)
	}

	// VerifyClaims returns the rich claim set incl. email/name.
	rawTok := iss.mint(t, valid)
	c, err := o.VerifyClaims(rawTok)
	if err != nil {
		t.Fatalf("VerifyClaims rejected valid token: %v", err)
	}
	if c.Subject != "user-1" || c.Email != "user-1@example.com" || c.Name != "User One" || len(c.Scopes) != 2 {
		t.Fatalf("unexpected claims: %+v", c)
	}

	// aud as array form also accepted
	arr := map[string]any{
		"iss": iss.srv.URL, "sub": "user-1", "aud": []string{"other", "my-client"},
		"exp": time.Now().Add(time.Hour).Unix(),
	}
	if _, err := o.Authenticate(bearerReq(iss.mint(t, arr))); err != nil {
		t.Fatalf("array aud rejected: %v", err)
	}

	cases := map[string]map[string]any{
		"expired":        {"iss": iss.srv.URL, "sub": "u", "aud": "my-client", "exp": time.Now().Add(-time.Hour).Unix()},
		"wrong issuer":   {"iss": "https://evil.example", "sub": "u", "aud": "my-client", "exp": time.Now().Add(time.Hour).Unix()},
		"wrong audience": {"iss": iss.srv.URL, "sub": "u", "aud": "not-me", "exp": time.Now().Add(time.Hour).Unix()},
	}
	for name, c := range cases {
		if _, err := o.Authenticate(bearerReq(iss.mint(t, c))); err == nil {
			t.Errorf("%s token accepted", name)
		}
		if _, err := o.VerifyClaims(iss.mint(t, c)); err == nil {
			t.Errorf("VerifyClaims accepted %s token", name)
		}
	}

	// tampered payload must fail signature verification
	tok := iss.mint(t, valid)
	tampered := fmt.Sprintf("%s.%s.%s",
		tok[:len(tok)/3],
		base64.RawURLEncoding.EncodeToString([]byte(`{"iss":"x","sub":"hacker"}`)),
		"AAAA")
	if _, err := o.Authenticate(bearerReq(tampered)); err == nil {
		t.Error("tampered token accepted")
	}
	if _, err := o.VerifyClaims(tampered); err == nil {
		t.Error("VerifyClaims accepted tampered token")
	}

	// no credentials → ErrNoCredentials so Chain falls through
	if _, err := o.Authenticate(httptest.NewRequest("GET", "/", nil)); err == nil {
		t.Error("missing bearer accepted")
	}
	if _, err := o.VerifyClaims(""); err == nil {
		t.Error("VerifyClaims accepted empty token")
	}
}

// --- API keys ---

func TestKeyStoreMintAuthenticateRevoke(t *testing.T) {
	path := filepath.Join(t.TempDir(), "apikeys.json")
	ks, err := OpenKeyStore(path)
	if err != nil {
		t.Fatal(err)
	}
	key, rec, err := ks.Mint("ci", []string{"nodes:read"})
	if err != nil {
		t.Fatal(err)
	}

	a := APIKeyAuth{Keys: ks}
	r := httptest.NewRequest("GET", "/api/nodes", nil)
	r.Header.Set("X-API-Key", key)
	p, err := a.Authenticate(r)
	if err != nil {
		t.Fatalf("minted key rejected: %v", err)
	}
	if p.Subject != rec.ID || p.Kind != KindAPIKey || !p.HasScope("nodes:read") {
		t.Fatalf("unexpected principal: %+v", p)
	}

	// persists across reopen
	ks2, err := OpenKeyStore(path)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := (APIKeyAuth{Keys: ks2}).Authenticate(r); err != nil {
		t.Fatalf("key lost after reopen: %v", err)
	}

	// invalid key rejected
	bad := httptest.NewRequest("GET", "/", nil)
	bad.Header.Set("X-API-Key", "gck_not_a_real_key")
	if _, err := a.Authenticate(bad); err == nil {
		t.Error("invalid key accepted")
	}

	// revoked key rejected
	if err := ks.Revoke(rec.ID); err != nil {
		t.Fatal(err)
	}
	if _, err := a.Authenticate(r); err == nil {
		t.Error("revoked key accepted")
	}
	if err := ks.Revoke(rec.ID); err == nil {
		t.Error("revoking twice should fail")
	}
}

func TestChainBearerThenAPIKey(t *testing.T) {
	iss := newFakeIssuer(t)
	ks, err := OpenKeyStore(filepath.Join(t.TempDir(), "k.json"))
	if err != nil {
		t.Fatal(err)
	}
	key, _, err := ks.Mint("integration", nil)
	if err != nil {
		t.Fatal(err)
	}
	chain := Chain{NewOIDC(iss.srv.URL, ""), APIKeyAuth{Keys: ks}}

	r := httptest.NewRequest("GET", "/", nil)
	r.Header.Set("X-API-Key", key)
	p, err := chain.Authenticate(r)
	if err != nil {
		t.Fatalf("api key via chain rejected: %v", err)
	}
	if p.Kind != KindAPIKey {
		t.Fatalf("expected apikey principal, got %+v", p)
	}

	tok := iss.mint(t, map[string]any{"iss": iss.srv.URL, "sub": "u1", "exp": time.Now().Add(time.Hour).Unix()})
	p, err = chain.Authenticate(bearerReq(tok))
	if err != nil {
		t.Fatalf("bearer via chain rejected: %v", err)
	}
	if p.Kind != KindOIDC {
		t.Fatalf("expected oidc principal, got %+v", p)
	}

	if _, err := chain.Authenticate(httptest.NewRequest("GET", "/", nil)); err == nil {
		t.Error("credential-less request accepted by chain")
	}
}

// --- middleware ---

func TestBearerMiddleware(t *testing.T) {
	iss := newFakeIssuer(t)
	tok := iss.mint(t, map[string]any{
		"iss": iss.srv.URL, "sub": "u", "aud": "client",
		"exp": time.Now().Add(time.Hour).Unix(),
	})
	a := NewOIDC(iss.srv.URL, "client")
	hits := 0
	handler := BearerMiddleware(a, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		hits++
		w.WriteHeader(http.StatusOK)
	}))

	// valid token → passes through
	w := httptest.NewRecorder()
	handler.ServeHTTP(w, bearerReq(tok))
	if w.Code != http.StatusOK || hits != 1 {
		t.Fatalf("valid token: code=%d hits=%d", w.Code, hits)
	}

	// missing token → 401
	w = httptest.NewRecorder()
	handler.ServeHTTP(w, httptest.NewRequest("GET", "/", nil))
	if w.Code != http.StatusUnauthorized || hits != 1 {
		t.Fatalf("missing token: code=%d hits=%d", w.Code, hits)
	}
}

func TestWriteMiddlewarePassesReads(t *testing.T) {
	a := NewOIDC("https://issuer.example", "")
	hits := 0
	handler := WriteMiddleware(a, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { hits++ }))

	for _, m := range []string{http.MethodGet, http.MethodHead, http.MethodOptions} {
		w := httptest.NewRecorder()
		handler.ServeHTTP(w, httptest.NewRequest(m, "/", nil))
		if w.Code != http.StatusOK {
			t.Errorf("%s: code=%d, want 200", m, w.Code)
		}
	}
}