package auth

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

func TestAudienceUnmarshal(t *testing.T) {
	var a1 Audience
	if err := json.Unmarshal([]byte(`"single-aud"`), &a1); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !a1.Contains("single-aud") || a1.Contains("other") {
		t.Fatalf("unexpected audience content: %v", a1)
	}

	var a2 Audience
	if err := json.Unmarshal([]byte(`["aud1", "aud2"]`), &a2); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if !a2.Contains("aud1") || !a2.Contains("aud2") || a2.Contains("aud3") {
		t.Fatalf("unexpected audience content: %v", a2)
	}
}

func TestVerifyToken(t *testing.T) {
	privKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("failed to generate RSA key: %v", err)
	}

	auth := &OidcAuth{
		issuer: "https://test.issuer",
		client: "test-client",
		keys: map[string]*rsa.PublicKey{
			"key-1": &privKey.PublicKey,
		},
		fetched: time.Now(),
	}

	// Helper to mint test token
	mintToken := func(kid string, claims IDClaims) string {
		headerJSON, _ := json.Marshal(map[string]string{
			"kid": kid,
			"alg": "RS256",
		})
		claimsJSON, _ := json.Marshal(claims)

		headerB64 := base64.RawURLEncoding.EncodeToString(headerJSON)
		payloadB64 := base64.RawURLEncoding.EncodeToString(claimsJSON)

		sum := sha256.Sum256([]byte(headerB64 + "." + payloadB64))
		sig, _ := rsa.SignPKCS1v15(rand.Reader, privKey, crypto.SHA256, sum[:])
		sigB64 := base64.RawURLEncoding.EncodeToString(sig)

		return headerB64 + "." + payloadB64 + "." + sigB64
	}

	// Valid token
	validToken := mintToken("key-1", IDClaims{
		Iss: "https://test.issuer",
		Sub: "user-123",
		Aud: Audience{"test-client"},
		Exp: time.Now().Add(time.Hour).Unix(),
	})

	claims, err := auth.VerifyToken(context.Background(), validToken)
	if err != nil {
		t.Fatalf("expected valid token to pass, got: %v", err)
	}
	if claims.Sub != "user-123" {
		t.Fatalf("expected sub user-123, got %s", claims.Sub)
	}

	// Expired token
	expiredToken := mintToken("key-1", IDClaims{
		Iss: "https://test.issuer",
		Sub: "user-123",
		Aud: Audience{"test-client"},
		Exp: time.Now().Add(-time.Hour).Unix(),
	})
	if _, err := auth.VerifyToken(context.Background(), expiredToken); err == nil {
		t.Fatalf("expected expired token to fail")
	}

	// Wrong audience
	wrongAudToken := mintToken("key-1", IDClaims{
		Iss: "https://test.issuer",
		Sub: "user-123",
		Aud: Audience{"other-client"},
		Exp: time.Now().Add(time.Hour).Unix(),
	})
	if _, err := auth.VerifyToken(context.Background(), wrongAudToken); err == nil {
		t.Fatalf("expected wrong audience to fail")
	}

	// Wrong issuer
	wrongIssToken := mintToken("key-1", IDClaims{
		Iss: "https://wrong.issuer",
		Sub: "user-123",
		Aud: Audience{"test-client"},
		Exp: time.Now().Add(time.Hour).Unix(),
	})
	if _, err := auth.VerifyToken(context.Background(), wrongIssToken); err == nil {
		t.Fatalf("expected wrong issuer to fail")
	}
}

func TestRequireAuthMiddleware(t *testing.T) {
	auth := &OidcAuth{
		issuer:     "https://test.issuer",
		client:     "test-client",
		ingestKey:  "secret-key",
		agentToken: "agent-secret",
	}

	handler := auth.RequireAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	}))

	// 1. Unauthenticated request -> 401
	req := httptest.NewRequest("GET", "/api/history", nil)
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("expected 401 for unauthenticated request, got %d", rec.Code)
	}

	// 2. Ingest key in header -> 200
	reqKey := httptest.NewRequest("GET", "/api/history", nil)
	reqKey.Header.Set("X-Otel-Ingest-Key", "secret-key")
	recKey := httptest.NewRecorder()
	handler.ServeHTTP(recKey, reqKey)
	if recKey.Code != http.StatusOK {
		t.Fatalf("expected 200 with valid ingest key, got %d", recKey.Code)
	}

	// 3. Dedicated service token in a header -> 200
	reqAgent := httptest.NewRequest("GET", "/api/history", nil)
	reqAgent.Header.Set("X-Otel-Agent-Token", "agent-secret")
	recAgent := httptest.NewRecorder()
	handler.ServeHTTP(recAgent, reqAgent)
	if recAgent.Code != http.StatusOK {
		t.Fatalf("expected 200 with valid agent token, got %d", recAgent.Code)
	}

	// 4. X-Auth-User from upstream ingress -> 200
	reqAuthUser := httptest.NewRequest("GET", "/api/history", nil)
	reqAuthUser.Header.Set("X-Auth-User", "mats@tornberg.me")
	recAuthUser := httptest.NewRecorder()
	handler.ServeHTTP(recAuthUser, reqAuthUser)
	if recAuthUser.Code != http.StatusOK {
		t.Fatalf("expected 200 with X-Auth-User, got %d", recAuthUser.Code)
	}
}