package auth

import (
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

// withResolver sets the deployment resource-identifier resolver for the duration
// of a test and restores the previous value afterward.
func withResolver(t *testing.T, resource string) {
	t.Helper()
	prev := resourceIdentifierResolver
	if resource == "" {
		SetResourceIdentifierResolver(nil)
	} else {
		SetResourceIdentifierResolver(func() string { return resource })
	}
	t.Cleanup(func() { resourceIdentifierResolver = prev })
}

func TestTokenAudienceMatches(t *testing.T) {
	cases := []struct {
		name string
		aud  jwt.ClaimStrings
		want string
		ok   bool
	}{
		{"unconfigured deployment allows anything", jwt.ClaimStrings{"https://a"}, "", true},
		{"no audience is grandfathered", nil, "https://prod", true},
		{"matching audience allowed", jwt.ClaimStrings{"https://prod"}, "https://prod", true},
		{"mismatched audience rejected", jwt.ClaimStrings{"https://stage"}, "https://prod", false},
		{"multi-audience with a match allowed", jwt.ClaimStrings{"https://x", "https://prod"}, "https://prod", true},
	}
	for _, tc := range cases {
		if got := tokenAudienceMatches(tc.aud, tc.want); got != tc.ok {
			t.Fatalf("%s: tokenAudienceMatches=%v want %v", tc.name, got, tc.ok)
		}
	}
}

// TestResourceBindingRejectsCrossEnvironmentToken is the core contract: a token
// minted for one deployment (staging) must be rejected by another (prod), even
// though both share the same signing secret in this test process.
func TestResourceBindingRejectsCrossEnvironmentToken(t *testing.T) {
	withResolver(t, "https://stage.example.com")
	token, _, err := GenerateToken(UserSession{ID: "u1"})
	if err != nil {
		t.Fatalf("mint: %v", err)
	}

	// Same deployment accepts it.
	if _, err := ValidateToken(token); err != nil {
		t.Fatalf("same-resource validation should pass, got %v", err)
	}

	// A different deployment (prod) rejects it.
	SetResourceIdentifierResolver(func() string { return "https://prod.example.com" })
	if _, err := ValidateToken(token); err == nil {
		t.Fatal("cross-resource validation should be rejected")
	} else if !strings.Contains(err.Error(), "not valid for this resource") {
		t.Fatalf("unexpected error: %v", err)
	}
}

// TestResourceBindingGrandfathersUnstampedToken proves the rollout is safe: a
// token minted before binding shipped (no audience) is still accepted once a
// deployment turns binding on — so enabling the feature never logs anyone out.
func TestResourceBindingGrandfathersUnstampedToken(t *testing.T) {
	withResolver(t, "") // binding off at mint → no aud stamped
	legacy, _, err := GenerateToken(UserSession{ID: "u1"})
	if err != nil {
		t.Fatalf("mint: %v", err)
	}
	if claims, _ := ValidateToken(legacy); claims != nil && len(claims.Audience) != 0 {
		t.Fatalf("expected no audience on legacy token, got %v", claims.Audience)
	}

	// Deployment now enforces a resource; the unstamped legacy token still passes.
	SetResourceIdentifierResolver(func() string { return "https://prod.example.com" })
	if _, err := ValidateToken(legacy); err != nil {
		t.Fatalf("unstamped legacy token should be grandfathered, got %v", err)
	}
}

// TestResourceBindingStampsAudienceAtMint confirms the aud claim is written when a
// resource is configured.
func TestResourceBindingStampsAudienceAtMint(t *testing.T) {
	withResolver(t, "https://prod.example.com")
	_, claims, err := GenerateToken(UserSession{ID: "u1"})
	if err != nil {
		t.Fatalf("mint: %v", err)
	}
	if len(claims.Audience) != 1 || claims.Audience[0] != "https://prod.example.com" {
		t.Fatalf("aud = %v, want [https://prod.example.com]", claims.Audience)
	}
}

func TestResourceBoundTokenRequiresOneExactAudience(t *testing.T) {
	const origin = "https://nexus.example"
	const mcpResource = origin + "/mcp"
	withResolver(t, origin)

	token, claims, err := GenerateTokenForSessionWithAudience(UserSession{ID: "u1"}, "session-1", time.Now(), DefaultSessionLifetime, mcpResource)
	if err != nil {
		t.Fatalf("GenerateTokenForSessionWithAudience: %v", err)
	}
	if len(claims.Audience) != 1 || claims.Audience[0] != mcpResource {
		t.Fatalf("aud = %v, want [%s]", claims.Audience, mcpResource)
	}
	if _, err := ValidateToken(token); err == nil {
		t.Fatal("ordinary host validator accepted an MCP-resource token")
	}
	if _, err := ValidateTokenForAudience(token, mcpResource); err != nil {
		t.Fatalf("exact MCP validator rejected MCP token: %v", err)
	}
	if _, err := ValidateTokenForAudience(token, origin); err == nil {
		t.Fatal("exact validator accepted MCP token for deployment origin")
	}

	legacy, _, err := GenerateTokenForSession(UserSession{ID: "u1"}, "session-2", time.Now(), DefaultSessionLifetime)
	if err != nil {
		t.Fatalf("GenerateTokenForSession: %v", err)
	}
	if _, err := ValidateTokenForAudience(legacy, mcpResource); err == nil {
		t.Fatal("exact validator accepted origin-audience token for MCP")
	}
}