package httpapi

import (
	"bytes"
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
	"github.com/golang-jwt/jwt/v5"
)

func TestOIDCJITFullGatewayUserFlowAndSecurityBoundary(t *testing.T) {
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run OIDC JIT HTTP integration tests")
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	applyMigration(t, ctx, databaseURL)
	db, err := store.Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()

	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("generate test signing key: %v", err)
	}
	var issuer string
	var gatewayBaseURL string
	var expectedPKCEChallenge string
	validSubject := ""
	issuerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/.well-known/openid-configuration":
			_ = json.NewEncoder(w).Encode(map[string]any{
				"issuer": issuer, "jwks_uri": issuer + "/jwks",
				"authorization_endpoint": issuer + "/authorize", "token_endpoint": issuer + "/token",
				"revocation_endpoint": issuer + "/revoke", "end_session_endpoint": issuer + "/logout",
			})
		case "/jwks":
			_ = json.NewEncoder(w).Encode(map[string]any{"keys": []any{oidcJITECJWK("jit-key", &key.PublicKey)}})
		case "/authorize":
			if r.URL.Query().Get("response_type") != "code" || r.URL.Query().Get("client_id") != "gateway-public-test" ||
				r.URL.Query().Get("code_challenge_method") != "S256" || r.URL.Query().Get("nonce") == "" {
				http.Error(w, "invalid authorization request", http.StatusBadRequest)
				return
			}
			expectedPKCEChallenge = r.URL.Query().Get("code_challenge")
			callback := gatewayBaseURL + "/api/v1/auth/oidc/callback?code=test-code&state=" + url.QueryEscape(r.URL.Query().Get("state")) +
				"&nonce=" + url.QueryEscape(r.URL.Query().Get("nonce"))
			http.Redirect(w, r, callback, http.StatusSeeOther)
		case "/token":
			if err := r.ParseForm(); err != nil || r.Form.Get("grant_type") != "authorization_code" ||
				r.Form.Get("client_id") != "gateway-public-test" || r.Form.Get("client_secret") != "" || r.Header.Get("Authorization") != "" {
				http.Error(w, "invalid token request", http.StatusBadRequest)
				return
			}
			digest := sha256.Sum256([]byte(r.Form.Get("code_verifier")))
			if base64.RawURLEncoding.EncodeToString(digest[:]) != expectedPKCEChallenge {
				http.Error(w, "invalid PKCE verifier", http.StatusBadRequest)
				return
			}
			accessToken := signedOIDCJITToken(t, key, issuer, validSubject, nil)
			idToken := signedOIDCJITToken(t, key, issuer, validSubject, func(claims jwt.MapClaims) {
				claims["aud"] = "gateway-public-test"
				// The nonce is retained from the authorization request by this test issuer.
				claims["nonce"] = currentOIDCTestNonce
			})
			_ = json.NewEncoder(w).Encode(map[string]any{
				"access_token": accessToken, "refresh_token": "opaque-test-refresh-token", "id_token": idToken,
				"token_type": "Bearer", "expires_in": 300,
			})
		case "/revoke":
			w.WriteHeader(http.StatusOK)
		case "/logout":
			http.Redirect(w, r, r.URL.Query().Get("post_logout_redirect_uri"), http.StatusSeeOther)
		default:
			http.NotFound(w, r)
		}
	}))
	defer issuerServer.Close()
	issuer = issuerServer.URL

	suffix := time.Now().UTC().Format("20060102150405.000000000")
	validSubject = "platform-http-jit-" + suffix
	rejectedSubjects := []string{
		"platform-http-scope-" + suffix,
		"platform-http-role-" + suffix,
		"platform-http-tenant-" + suffix,
		"platform-http-disabled-jit-" + suffix,
		"platform-http-missing-tenant-" + suffix,
	}
	allSubjects := append([]string{validSubject}, rejectedSubjects...)
	t.Cleanup(func() {
		_, _ = db.Pool().Exec(context.Background(), `
DELETE FROM gateway_audit_logs
WHERE target_id IN (SELECT id::text FROM gateway_users WHERE source = 'oidc' AND external_user_id = ANY($1::text[]));
DELETE FROM gateway_users WHERE source = 'oidc' AND external_user_id = ANY($1::text[]);`, allSubjects)
	})

	baseConfig := config.Config{
		AppEnv:                          "test",
		HTTPAddr:                        ":0",
		DatabaseURL:                     databaseURL,
		IdentityMode:                    "hybrid",
		JWTSecret:                       "test-only-jwt-secret",
		OIDCEnabled:                     true,
		OIDCIssuer:                      issuer,
		OIDCAudience:                    "gateway-api",
		OIDCTenantID:                    "auth-center-test-tenant",
		OIDCRolePrefix:                  "gateway.",
		OIDCRequiredScopes:              []string{"gateway.access"},
		OIDCJWKSCacheTTLSeconds:         60,
		OIDCAcceptLegacyHS256:           true,
		OIDCJITProvisioningEnabled:      true,
		OIDCGatewayTenantKey:            "default",
		OIDCBrowserSessionEnabled:       true,
		OIDCSessionCookieSecure:         false,
		OIDCClientID:                    "gateway-public-test",
		OIDCRedirectURI:                 "http://localhost/api/v1/auth/oidc/callback",
		OIDCPostLogoutRedirectURI:       "http://localhost:5178/",
		OIDCSessionEncryptionKey:        base64.RawStdEncoding.EncodeToString(bytes.Repeat([]byte{8}, 32)),
		OIDCSessionIdleTTLSeconds:       1800,
		OIDCSessionAbsoluteTTLSeconds:   28800,
		OIDCSessionRefreshBeforeSeconds: 60,
		LocalGeneratedStorageDir:        t.TempDir(),
		LocalUploadedStorageDir:         t.TempDir(),
		LocalTempAssetTTLHours:          1,
		CORSAllowedOrigin:               "http://localhost:5178",
		TaskProgressCallbackEnabled:     false,
	}
	server := httptest.NewServer(NewServerWithContext(ctx, baseConfig, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
	defer server.Close()
	gatewayBaseURL = server.URL

	validToken := signedOIDCJITToken(t, key, issuer, validSubject, nil)
	var me auth.User
	doOIDCJITJSON(t, server.URL, http.MethodGet, "/api/v1/me", validToken, nil, http.StatusOK, &me)
	if me.ID != validSubject || me.Source != "oidc" || me.GatewayUserID == "" || me.GatewayTenantID == "" || me.TenantKey != "default" || me.UserGroupID == "" {
		t.Fatalf("OIDC /me did not include the local Gateway projection")
	}
	sessionCookie := createOIDCBFFSessionCookie(t, server.URL)
	request, err := http.NewRequest(http.MethodGet, server.URL+"/api/v1/me", nil)
	if err != nil {
		t.Fatal(err)
	}
	request.AddCookie(sessionCookie)
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		t.Fatalf("execute cookie-authenticated /me: %v", err)
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		t.Fatalf("cookie-authenticated /me status = %d, want 200", response.StatusCode)
	}
	var cookieMe auth.User
	if err := json.NewDecoder(response.Body).Decode(&cookieMe); err != nil {
		t.Fatalf("decode cookie-authenticated /me: %v", err)
	}
	if cookieMe.GatewayUserID != me.GatewayUserID || cookieMe.ID != me.ID {
		t.Fatalf("new-tab cookie resolved a different Gateway user: %#v", cookieMe)
	}
	var automaticallyCreatedAPIKeys int
	if err := db.Pool().QueryRow(ctx, `SELECT count(*) FROM gateway_api_keys WHERE gateway_user_id = $1::uuid`, me.GatewayUserID).Scan(&automaticallyCreatedAPIKeys); err != nil {
		t.Fatalf("count pre-created API keys: %v", err)
	}
	if automaticallyCreatedAPIKeys != 0 {
		t.Fatalf("OIDC JIT created %d API keys before explicit user action", automaticallyCreatedAPIKeys)
	}

	for _, path := range []string{
		"/api/workspace/user-groups",
		"/api/workspace/wallet",
		"/api/workspace/tasks",
		"/api/v1/api-keys",
	} {
		doOIDCJITJSON(t, server.URL, http.MethodGet, path, validToken, nil, http.StatusOK, nil)
	}

	var createdKey struct {
		Secret string `json:"secret"`
		APIKey struct {
			ID string `json:"id"`
		} `json:"apiKey"`
	}
	doOIDCJITJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", validToken, map[string]any{"name": "OIDC JIT integration key"}, http.StatusCreated, &createdKey)
	if createdKey.Secret == "" || createdKey.APIKey.ID == "" {
		t.Fatal("OIDC user API Key creation returned incomplete data")
	}
	doOIDCJITJSON(t, server.URL, http.MethodGet, "/api/v1/api-keys", validToken, nil, http.StatusOK, nil)
	doOIDCJITJSON(t, server.URL, http.MethodDelete, "/api/v1/api-keys/"+createdKey.APIKey.ID, validToken, nil, http.StatusNoContent, nil)

	var users struct {
		Items []store.GatewayUser `json:"items"`
	}
	doOIDCJITJSON(t, server.URL, http.MethodGet, "/api/admin/users", validToken, nil, http.StatusOK, &users)
	foundOIDCUser := false
	for _, user := range users.Items {
		if user.ID == me.GatewayUserID {
			foundOIDCUser = user.Source == "oidc" && user.ExternalUserID == validSubject
			break
		}
	}
	if !foundOIDCUser {
		t.Fatal("admin user list did not expose the OIDC Gateway projection")
	}

	negativeTokens := []struct {
		subject string
		mutate  func(jwt.MapClaims)
	}{
		{rejectedSubjects[0], func(claims jwt.MapClaims) { claims["scope"] = "openid" }},
		{rejectedSubjects[1], func(claims jwt.MapClaims) { claims["roles"] = []string{"other.admin"} }},
		{rejectedSubjects[2], func(claims jwt.MapClaims) { claims["tid"] = "wrong-tenant" }},
	}
	for _, negative := range negativeTokens {
		token := signedOIDCJITToken(t, key, issuer, negative.subject, negative.mutate)
		doOIDCJITJSON(t, server.URL, http.MethodGet, "/api/v1/me", token, nil, http.StatusUnauthorized, nil)
	}
	var rejectedWrites int
	if err := db.Pool().QueryRow(ctx, `SELECT count(*) FROM gateway_users WHERE source = 'oidc' AND external_user_id = ANY($1::text[])`, rejectedSubjects[:3]).Scan(&rejectedWrites); err != nil {
		t.Fatalf("count rejected OIDC writes: %v", err)
	}
	if rejectedWrites != 0 {
		t.Fatalf("rejected OIDC tokens created %d local users", rejectedWrites)
	}

	disabledJITConfig := baseConfig
	disabledJITConfig.OIDCJITProvisioningEnabled = false
	disabledJITServer := httptest.NewServer(NewServerWithContext(ctx, disabledJITConfig, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
	defer disabledJITServer.Close()
	assertOIDCJITError(t, disabledJITServer.URL, signedOIDCJITToken(t, key, issuer, rejectedSubjects[3], nil), http.StatusForbidden, errorCodeGatewayUserNotProvisioned)

	missingTenantConfig := baseConfig
	missingTenantConfig.OIDCGatewayTenantKey = "missing-tenant-" + suffix
	missingTenantServer := httptest.NewServer(NewServerWithContext(ctx, missingTenantConfig, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
	defer missingTenantServer.Close()
	assertOIDCJITError(t, missingTenantServer.URL, signedOIDCJITToken(t, key, issuer, rejectedSubjects[4], nil), http.StatusServiceUnavailable, errorCodeGatewayTenantUnavailable)

	if _, err := db.Pool().Exec(ctx, `UPDATE gateway_users SET status = 'disabled' WHERE id = $1::uuid`, me.GatewayUserID); err != nil {
		t.Fatalf("disable projected user: %v", err)
	}
	assertOIDCJITError(t, server.URL, validToken, http.StatusForbidden, errorCodeGatewayUserDisabled)
	if _, err := db.Pool().Exec(ctx, `UPDATE gateway_users SET status = 'active' WHERE id = $1::uuid`, me.GatewayUserID); err != nil {
		t.Fatalf("restore projected user for delete test: %v", err)
	}
	if err := db.DeleteGatewayUser(ctx, me.GatewayUserID); err != nil {
		t.Fatalf("delete projected user: %v", err)
	}
	assertOIDCJITError(t, server.URL, validToken, http.StatusForbidden, errorCodeGatewayUserDisabled)
}

var currentOIDCTestNonce string

func createOIDCBFFSessionCookie(t *testing.T, baseURL string) *http.Cookie {
	t.Helper()
	jar, err := cookiejar.New(nil)
	if err != nil {
		t.Fatal(err)
	}
	client := &http.Client{Jar: jar, CheckRedirect: func(request *http.Request, via []*http.Request) error {
		if request.URL.Path == "/api/v1/auth/oidc/callback" {
			currentOIDCTestNonce = request.URL.Query().Get("nonce")
		}
		if len(via) > 10 {
			return errors.New("too many redirects")
		}
		return nil
	}}
	response, err := client.Get(baseURL + "/api/v1/auth/oidc/login?returnTo=%2Fapi%2Fv1%2Fme")
	if err != nil {
		t.Fatalf("complete OIDC BFF login: %v", err)
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(response.Body)
		t.Fatalf("OIDC BFF login status = %d, want 200: %s", response.StatusCode, body)
	}
	parsedBaseURL, _ := url.Parse(baseURL)
	for _, cookie := range jar.Cookies(parsedBaseURL) {
		if cookie.Name == auth.OIDCSessionCookieName {
			if strings.Count(cookie.Value, ".") == 2 {
				t.Fatal("browser session cookie contains a JWT instead of an opaque session ID")
			}
			return cookie
		}
	}
	t.Fatal("OIDC browser session cookie was not returned")
	return nil
}

func assertOIDCJITError(t *testing.T, baseURL string, token string, expectedStatus int, expectedCode string) {
	t.Helper()
	var envelope struct {
		Error struct {
			Code string `json:"code"`
		} `json:"error"`
	}
	doOIDCJITJSON(t, baseURL, http.MethodGet, "/api/v1/me", token, nil, expectedStatus, &envelope)
	if envelope.Error.Code != expectedCode {
		t.Fatalf("error code = %q, want %q", envelope.Error.Code, expectedCode)
	}
}

func doOIDCJITJSON(t *testing.T, baseURL string, method string, path string, token string, payload any, expectedStatus int, output any) {
	t.Helper()
	var body io.Reader
	if payload != nil {
		raw, err := json.Marshal(payload)
		if err != nil {
			t.Fatalf("marshal OIDC JIT request: %v", err)
		}
		body = bytes.NewReader(raw)
	}
	request, err := http.NewRequest(method, baseURL+path, body)
	if err != nil {
		t.Fatalf("build %s %s request: %v", method, path, err)
	}
	request.Header.Set("Authorization", "Bearer "+token)
	if payload != nil {
		request.Header.Set("Content-Type", "application/json")
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		t.Fatalf("execute %s %s: %v", method, path, err)
	}
	defer response.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(response.Body, 1<<20))
	if err != nil {
		t.Fatalf("read %s %s response: %v", method, path, err)
	}
	if response.StatusCode != expectedStatus {
		t.Fatalf("%s %s status=%d, want=%d", method, path, response.StatusCode, expectedStatus)
	}
	if output != nil && len(raw) > 0 {
		if err := json.Unmarshal(raw, output); err != nil {
			t.Fatalf("decode %s %s response: %v", method, path, err)
		}
	}
}

func signedOIDCJITToken(t *testing.T, key *ecdsa.PrivateKey, issuer string, subject string, mutate func(jwt.MapClaims)) string {
	t.Helper()
	now := time.Now()
	claims := jwt.MapClaims{
		"iss": issuer, "aud": "gateway-api", "sub": subject, "tid": "auth-center-test-tenant",
		"preferred_username": "oidc-jit-acceptance", "roles": []string{"gateway.admin"},
		"scope": "openid gateway.access", "iat": now.Unix(), "nbf": now.Add(-time.Second).Unix(),
		"exp": now.Add(time.Hour).Unix(),
	}
	if mutate != nil {
		mutate(claims)
	}
	token := jwt.NewWithClaims(jwt.SigningMethodES256, claims)
	token.Header["kid"] = "jit-key"
	raw, err := token.SignedString(key)
	if err != nil {
		t.Fatalf("sign OIDC JIT test token: %v", err)
	}
	return raw
}

func oidcJITECJWK(kid string, key *ecdsa.PublicKey) map[string]any {
	return map[string]any{
		"kid": kid,
		"kty": "EC",
		"use": "sig",
		"alg": "ES256",
		"crv": "P-256",
		"x":   base64.RawURLEncoding.EncodeToString(key.X.FillBytes(make([]byte, 32))),
		"y":   base64.RawURLEncoding.EncodeToString(key.Y.FillBytes(make([]byte, 32))),
	}
}