package middleware

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/sandbox0-ai/sandbox0/pkg/gateway/apikey"
	"github.com/sandbox0-ai/sandbox0/pkg/gateway/authn"
	"github.com/sandbox0-ai/sandbox0/pkg/gateway/spec"
	"github.com/sandbox0-ai/sandbox0/pkg/internalauth"
	"go.uber.org/zap"
)

type staticAPIKeyValidator struct {
	key *apikey.APIKey
}

func (v staticAPIKeyValidator) ValidateAPIKey(context.Context, string) (*apikey.APIKey, error) {
	return v.key, nil
}

func TestAuthMiddlewareAuthenticateReturnsErrorEnvelope(t *testing.T) {
	gin.SetMode(gin.TestMode)

	middleware := NewAuthMiddleware(nil, "test-secret", nil, zap.NewNop())
	router := gin.New()
	router.Use(middleware.Authenticate())
	router.GET("/protected", func(c *gin.Context) {
		c.Status(http.StatusNoContent)
	})

	req := httptest.NewRequest(http.MethodGet, "/protected", nil)
	rec := httptest.NewRecorder()

	router.ServeHTTP(rec, req)

	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("status = %d, want %d", rec.Code, http.StatusUnauthorized)
	}
	var response spec.Response
	if err := json.Unmarshal(rec.Body.Bytes(), &response); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if response.Success || response.Error == nil {
		t.Fatalf("response = %+v, want error envelope", response)
	}
	if response.Error.Code != spec.CodeUnauthorized || response.Error.Message != "missing authorization header" {
		t.Fatalf("error = %+v, want unauthorized missing authorization header", response.Error)
	}
}

func TestAuthMiddleware_JWTAccessToken(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	issuer := authn.NewIssuer("regional-gateway", "test-secret", time.Minute, time.Hour)
	tokens, err := issuer.IssueTokenPair("user-1", "user@example.com", "User", false, []authn.TeamGrant{
		{TeamID: "team-1", TeamRole: "admin", HomeRegionID: "aws-us-east-1"},
		{TeamID: "team-2", TeamRole: "developer", HomeRegionID: "aws-us-east-1"},
	})
	if err != nil {
		t.Fatalf("issue token pair: %v", err)
	}

	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer "+tokens.AccessToken)
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(nil, "test-secret", issuer, zap.NewNop())
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if authCtx.UserID != "user-1" || authCtx.TeamID != "" {
		t.Fatalf("unexpected auth context: user=%s team=%s", authCtx.UserID, authCtx.TeamID)
	}
}

func TestAuthMiddleware_JWTAccessTokenImplicitSingleTeamInDirectClusterMode(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	issuer := authn.NewIssuer("cluster-gateway", "test-secret", time.Minute, time.Hour)
	tokens, err := issuer.IssueTokenPair("user-1", "user@example.com", "User", false, []authn.TeamGrant{{TeamID: "team-1", TeamRole: "admin", HomeRegionID: "aws-us-east-1"}})
	if err != nil {
		t.Fatalf("issue token pair: %v", err)
	}

	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer "+tokens.AccessToken)
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(nil, "test-secret", issuer, zap.NewNop())
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if authCtx.UserID != "user-1" || authCtx.TeamID != "team-1" || authCtx.TeamRole != "admin" {
		t.Fatalf("unexpected auth context: %+v", authCtx)
	}
}

func TestAuthMiddleware_JWTRefreshTokenRejected(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	issuer := authn.NewIssuer("regional-gateway", "test-secret", time.Minute, time.Hour)
	tokens, err := issuer.IssueTokenPair("user-1", "user@example.com", "User", false, []authn.TeamGrant{{TeamID: "team-1", TeamRole: "admin", HomeRegionID: "aws-us-east-1"}})
	if err != nil {
		t.Fatalf("issue token pair: %v", err)
	}

	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer "+tokens.RefreshToken)
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(nil, "test-secret", issuer, zap.NewNop())
	if _, err := middleware.AuthenticateRequest(ctx); err == nil {
		t.Fatalf("expected refresh token to be rejected")
	}
}

func TestAuthMiddleware_APIKeyCarriesOwnerUserID(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	ownerID := "user-1"
	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer s0_test")
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(staticAPIKeyValidator{key: &apikey.APIKey{
		ID:        "key-1",
		TeamID:    "team-1",
		UserID:    &ownerID,
		CreatedBy: "creator-1",
		Roles:     []string{"developer"},
	}}, "test-secret", nil, zap.NewNop())
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if authCtx.UserID != ownerID {
		t.Fatalf("user id = %q, want %q", authCtx.UserID, ownerID)
	}
	if authCtx.TeamID != "team-1" || authCtx.APIKeyID != "key-1" {
		t.Fatalf("unexpected auth context: %+v", authCtx)
	}
}

func TestAuthMiddleware_APIKeyFallsBackToCreatorUserID(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer s0_test")
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(staticAPIKeyValidator{key: &apikey.APIKey{
		ID:        "key-1",
		TeamID:    "team-1",
		CreatedBy: "creator-1",
		Roles:     []string{"developer"},
	}}, "test-secret", nil, zap.NewNop())
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if authCtx.UserID != "creator-1" {
		t.Fatalf("user id = %q, want creator-1", authCtx.UserID)
	}
}

func TestAuthMiddleware_PlatformAPIKeySetsSystemAdminWithoutSelectedTeam(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	ownerID := "user-1"
	req := httptest.NewRequest("GET", "/internal/v1/metering/status", nil)
	req.Header.Set("Authorization", "Bearer s0_test")
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(staticAPIKeyValidator{key: &apikey.APIKey{
		ID:        "key-1",
		TeamID:    "team-1",
		UserID:    &ownerID,
		CreatedBy: "creator-1",
		Scope:     apikey.ScopePlatform,
		Roles:     []string{"viewer"},
	}}, "test-secret", nil, zap.NewNop())
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if !authCtx.IsSystemAdmin {
		t.Fatal("expected platform API key to set system admin")
	}
	if !authCtx.HasPermission(authn.PermTemplateDelete) || !authCtx.HasPermission("anything:anywhere") {
		t.Fatalf("expected wildcard permissions, got %v", authCtx.Permissions)
	}
	if authCtx.AuthMethod != authn.AuthMethodAPIKey || authCtx.APIKeyID != "key-1" {
		t.Fatalf("unexpected auth context: %+v", authCtx)
	}
	if authCtx.TeamID != "" {
		t.Fatalf("team id = %q, want empty without selected team", authCtx.TeamID)
	}
}

func TestAuthMiddleware_PlatformAPIKeyUsesSelectedTeamHeader(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	req := httptest.NewRequest("GET", "/api/v1/sandboxes/sbox-1", nil)
	req.Header.Set("Authorization", "Bearer s0_test")
	req.Header.Set(internalauth.TeamIDHeader, "team-target")
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(staticAPIKeyValidator{key: &apikey.APIKey{
		ID:        "key-1",
		TeamID:    "platform-owner-team",
		CreatedBy: "creator-1",
		Scope:     apikey.ScopePlatform,
	}}, "test-secret", nil, zap.NewNop())
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if !authCtx.IsSystemAdmin {
		t.Fatal("expected platform API key to set system admin")
	}
	if authCtx.TeamID != "team-target" {
		t.Fatalf("team id = %q, want selected team", authCtx.TeamID)
	}
}

func TestAuthMiddleware_TeamAPIKeyRejectsMismatchedTeamHeader(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	req := httptest.NewRequest("GET", "/api/v1/sandboxes/sbox-1", nil)
	req.Header.Set("Authorization", "Bearer s0_test")
	req.Header.Set(internalauth.TeamIDHeader, "team-other")
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(staticAPIKeyValidator{key: &apikey.APIKey{
		ID:        "key-1",
		TeamID:    "team-1",
		CreatedBy: "creator-1",
		Roles:     []string{"developer"},
	}}, "test-secret", nil, zap.NewNop())
	if _, err := middleware.AuthenticateRequest(ctx); err != ErrSelectedTeamForbidden {
		t.Fatalf("authenticate request error = %v, want %v", err, ErrSelectedTeamForbidden)
	}
}

func TestAuthMiddleware_JWTAccessTokenExplicitTeamHeader(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	homeRegionID := "aws-us-east-1"
	issuer := authn.NewIssuer("regional-gateway", "test-secret", time.Minute, time.Hour)
	tokens, err := issuer.IssueTokenPair("user-1", "user@example.com", "User", false, []authn.TeamGrant{{TeamID: "team-2", TeamRole: "developer", HomeRegionID: homeRegionID}})
	if err != nil {
		t.Fatalf("issue token pair: %v", err)
	}

	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer "+tokens.AccessToken)
	req.Header.Set(internalauth.TeamIDHeader, "team-2")
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(
		nil,
		"test-secret",
		issuer,
		zap.NewNop(),
		WithRequiredTeamRegionID(homeRegionID),
	)
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if authCtx.TeamID != "team-2" || authCtx.TeamRole != "developer" {
		t.Fatalf("unexpected auth context: team=%s role=%s", authCtx.TeamID, authCtx.TeamRole)
	}
	if !authCtx.HasPermission(authn.PermSandboxCreate) {
		t.Fatal("expected developer permissions to be populated from selected team membership")
	}
}

func TestAuthMiddleware_JWTAccessTokenExplicitTeamHeaderRequiredForTeamlessToken(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	issuer := authn.NewIssuer("regional-gateway", "test-secret", time.Minute, time.Hour)
	tokens, err := issuer.IssueTokenPair("user-1", "user@example.com", "User", false, nil)
	if err != nil {
		t.Fatalf("issue token pair: %v", err)
	}

	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer "+tokens.AccessToken)
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(nil, "test-secret", issuer, zap.NewNop())
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if authCtx.TeamID != "" || authCtx.TeamRole != "" {
		t.Fatalf("expected teamless auth context, got team=%q role=%q", authCtx.TeamID, authCtx.TeamRole)
	}
}

func TestAuthMiddleware_JWTAccessTokenExplicitTeamHeaderRejectsWrongRegion(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	issuer := authn.NewIssuer("regional-gateway", "test-secret", time.Minute, time.Hour)
	tokens, err := issuer.IssueTokenPair("user-1", "user@example.com", "User", false, []authn.TeamGrant{{TeamID: "team-2", TeamRole: "developer", HomeRegionID: "aws-us-west-2"}})
	if err != nil {
		t.Fatalf("issue token pair: %v", err)
	}

	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer "+tokens.AccessToken)
	req.Header.Set(internalauth.TeamIDHeader, "team-2")
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	middleware := NewAuthMiddleware(
		nil,
		"test-secret",
		issuer,
		zap.NewNop(),
		WithRequiredTeamRegionID("aws-us-east-1"),
	)
	if _, err := middleware.AuthenticateRequest(ctx); err != ErrSelectedTeamWrongRegion {
		t.Fatalf("authenticate request error = %v, want %v", err, ErrSelectedTeamWrongRegion)
	}
}

func TestAuthMiddleware_JWTAccessTokenWithVerifierOnlyIssuer(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	privateKeyPEM, publicKeyPEM, err := internalauth.GenerateEd25519KeyPair()
	if err != nil {
		t.Fatalf("generate key pair: %v", err)
	}
	privateKey, err := internalauth.LoadEd25519PrivateKey(privateKeyPEM)
	if err != nil {
		t.Fatalf("load private key: %v", err)
	}
	publicKey, err := internalauth.LoadEd25519PublicKey(publicKeyPEM)
	if err != nil {
		t.Fatalf("load public key: %v", err)
	}

	issuer := authn.NewIssuerWithEd25519("global-gateway", privateKey, time.Minute, time.Hour)
	tokens, err := issuer.IssueTokenPair("user-1", "user@example.com", "User", false, []authn.TeamGrant{{TeamID: "team-2", TeamRole: "developer", HomeRegionID: "aws-us-east-1"}})
	if err != nil {
		t.Fatalf("issue token pair: %v", err)
	}

	req := httptest.NewRequest("GET", "/api/v1/templates", nil)
	req.Header.Set("Authorization", "Bearer "+tokens.AccessToken)
	req.Header.Set(internalauth.TeamIDHeader, "team-2")
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req

	verifier := authn.NewVerifierWithEd25519("global-gateway", publicKey)
	middleware := NewAuthMiddleware(nil, "", verifier, zap.NewNop(), WithRequiredTeamRegionID("aws-us-east-1"))
	authCtx, err := middleware.AuthenticateRequest(ctx)
	if err != nil {
		t.Fatalf("authenticate request: %v", err)
	}
	if authCtx.UserID != "user-1" || authCtx.TeamID != "team-2" || authCtx.TeamRole != "developer" {
		t.Fatalf("unexpected auth context: %+v", authCtx)
	}
}

func TestAuthMiddleware_RequireJWTAuth(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	tests := []struct {
		name       string
		authCtx    *authn.AuthContext
		wantStatus int
	}{
		{
			name: "jwt user allowed",
			authCtx: &authn.AuthContext{
				AuthMethod: authn.AuthMethodJWT,
				UserID:     "user-1",
				TeamID:     "team-1",
			},
			wantStatus: http.StatusNoContent,
		},
		{
			name: "api key rejected",
			authCtx: &authn.AuthContext{
				AuthMethod: authn.AuthMethodAPIKey,
				TeamID:     "team-1",
			},
			wantStatus: http.StatusUnauthorized,
		},
		{
			name: "jwt without user rejected",
			authCtx: &authn.AuthContext{
				AuthMethod: authn.AuthMethodJWT,
				TeamID:     "team-1",
			},
			wantStatus: http.StatusUnauthorized,
		},
		{
			name:       "missing auth context rejected",
			authCtx:    nil,
			wantStatus: http.StatusUnauthorized,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			m := NewAuthMiddleware(nil, "test-secret", nil, zap.NewNop())
			engine := gin.New()
			engine.Use(func(c *gin.Context) {
				if tt.authCtx != nil {
					c.Set("auth_context", tt.authCtx)
				}
				c.Next()
			})
			engine.Use(m.RequireJWTAuth())
			engine.GET("/teams", func(c *gin.Context) {
				c.Status(http.StatusNoContent)
			})

			req := httptest.NewRequest(http.MethodGet, "/teams", nil)
			rec := httptest.NewRecorder()
			engine.ServeHTTP(rec, req)

			if rec.Code != tt.wantStatus {
				t.Fatalf("unexpected status: got %d want %d", rec.Code, tt.wantStatus)
			}
			if tt.wantStatus == http.StatusUnauthorized {
				var body spec.Response
				if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
					t.Fatalf("decode response: %v", err)
				}
				if body.Error == nil || body.Error.Message != "this API requires a user access token (human login); API keys are not supported" {
					t.Fatalf("unexpected error response: %+v", body.Error)
				}
			}
		})
	}
}

func TestAuthMiddleware_RequireSystemAdmin(t *testing.T) {
	t.Setenv("GIN_MODE", "release")

	tests := []struct {
		name       string
		authCtx    *authn.AuthContext
		wantStatus int
	}{
		{
			name: "system admin allowed",
			authCtx: &authn.AuthContext{
				AuthMethod:    authn.AuthMethodJWT,
				UserID:        "user-1",
				TeamID:        "team-1",
				IsSystemAdmin: true,
			},
			wantStatus: http.StatusNoContent,
		},
		{
			name: "platform api key allowed",
			authCtx: &authn.AuthContext{
				AuthMethod:    authn.AuthMethodAPIKey,
				UserID:        "user-1",
				TeamID:        "team-1",
				APIKeyID:      "key-1",
				IsSystemAdmin: true,
			},
			wantStatus: http.StatusNoContent,
		},
		{
			name: "team admin rejected",
			authCtx: &authn.AuthContext{
				AuthMethod: authn.AuthMethodJWT,
				UserID:     "user-1",
				TeamID:     "team-1",
				TeamRole:   "admin",
			},
			wantStatus: http.StatusForbidden,
		},
		{
			name:       "missing auth context rejected",
			authCtx:    nil,
			wantStatus: http.StatusUnauthorized,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			m := NewAuthMiddleware(nil, "test-secret", nil, zap.NewNop())
			engine := gin.New()
			engine.Use(func(c *gin.Context) {
				if tt.authCtx != nil {
					c.Set("auth_context", tt.authCtx)
				}
				c.Next()
			})
			engine.Use(m.RequireSystemAdmin())
			engine.GET("/internal/v1/metering/status", func(c *gin.Context) {
				c.Status(http.StatusNoContent)
			})

			req := httptest.NewRequest(http.MethodGet, "/internal/v1/metering/status", nil)
			rec := httptest.NewRecorder()
			engine.ServeHTTP(rec, req)

			if rec.Code != tt.wantStatus {
				t.Fatalf("unexpected status: got %d want %d", rec.Code, tt.wantStatus)
			}
		})
	}
}
