package store

import (
	"context"
	"crypto/sha256"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

func TestSecurityEventWatermarkVerificationAndFallbackLifecycle(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 security event PostgreSQL integration tests")
	}
	ctx := context.Background()
	probe, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		t.Fatal(err)
	}
	var databaseName string
	if err := probe.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		probe.Close()
		t.Fatal(err)
	}
	probe.Close()
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to migrate non-test database %q", databaseName)
	}
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	issuer := "https://auth.test.example/ssf"
	subjectIssuer := "https://auth.test/issuer/shared"
	audience := "urn:easyai:ssf:receiver:" + uuid.NewString()
	streamID, tenantID, subject := uuid.NewString(), uuid.NewString(), uuid.NewString()
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE transmitter_issuer=$1`, issuer)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_receipts WHERE issuer=$1`, issuer)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_oidc_revocation_watermarks WHERE issuer=$1`, subjectIssuer)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_stream_state WHERE issuer=$1`, issuer)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_audit_logs WHERE actor_source='ssf' AND target_id=ANY($1)`, []string{shortSecurityEventHash(subject), shortSecurityEventHash(streamID)})
	})
	if err := db.EnsureSecurityEventStreamState(ctx, issuer, audience, streamID); err != nil {
		t.Fatal(err)
	}
	now := time.Now().UTC().Truncate(time.Second)
	if err := db.AdvanceSecurityEventStreamState(ctx, issuer, audience, now, 180*time.Second); err != nil {
		t.Fatalf("advance fresh stream state: %v", err)
	}
	evaluation, err := db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, "", tenantID, subject, now, now, 180*time.Second)
	if err != nil || !evaluation.RequireIntrospection || evaluation.Mode != "bootstrap" {
		t.Fatalf("bootstrap evaluation=%#v error=%v", evaluation, err)
	}

	confirmAt := func(label string, at time.Time) {
		t.Helper()
		stateHash := sha256.Sum256([]byte(label))
		if err := db.BeginSecurityEventVerification(ctx, issuer, audience, stateHash[:], at); err != nil {
			t.Fatal(err)
		}
		matched, err := db.ConfirmSecurityEventVerification(ctx, issuer, audience, streamID, uuid.NewString(), stateHash[:], at)
		if err != nil || !matched {
			t.Fatalf("confirm matched=%v error=%v", matched, err)
		}
	}
	inserted, err := db.ApplySecurityEventStreamUpdated(ctx, issuer, audience, uuid.NewString(), streamID, "enabled", "onboarding", now)
	if err != nil || !inserted {
		t.Fatalf("enable stream inserted=%v error=%v", inserted, err)
	}
	confirmAt("first-verification-state", now)
	confirmAt("second-verification-state", now.Add(time.Minute))
	evaluation, err = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, "", tenantID, subject, now, now.Add(time.Minute), 180*time.Second)
	if err != nil || !evaluation.RequireIntrospection || evaluation.Mode != "introspection_fallback" {
		t.Fatalf("bootstrap overlap evaluation=%#v error=%v", evaluation, err)
	}
	confirmAt("post-bootstrap-verification-state", now.Add(361*time.Second))
	evaluation, err = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, "", tenantID, subject, now, now.Add(361*time.Second), 180*time.Second)
	if err != nil || evaluation.RequireIntrospection || evaluation.Mode != "push_healthy" {
		t.Fatalf("healthy evaluation=%#v error=%v", evaluation, err)
	}
	streamUpdateJTI := uuid.NewString()
	inserted, err = db.ApplySecurityEventStreamUpdated(ctx, issuer, audience, streamUpdateJTI, streamID, "paused", "maintenance", now.Add(362*time.Second))
	if err != nil || !inserted {
		t.Fatalf("stream update inserted=%v error=%v", inserted, err)
	}
	inserted, err = db.ApplySecurityEventStreamUpdated(ctx, issuer, audience, streamUpdateJTI, streamID, "paused", "maintenance", now)
	if err != nil || inserted {
		t.Fatalf("duplicate stream update inserted=%v error=%v", inserted, err)
	}
	confirmAt("paused-verification-one", now.Add(363*time.Second))
	confirmAt("paused-verification-two", now.Add(364*time.Second))
	evaluation, err = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, "", tenantID, subject, now, now.Add(364*time.Second), 180*time.Second)
	if err != nil || !evaluation.RequireIntrospection || evaluation.Mode != "introspection_fallback" {
		t.Fatalf("stream update fallback=%#v error=%v", evaluation, err)
	}
	_, _ = db.pool.Exec(ctx, `UPDATE gateway_security_event_stream_state SET stream_status='enabled',mode='push_healthy',fallback_since=NULL,
		consecutive_verifications=0,last_verification_at=$3 WHERE issuer=$1 AND audience=$2`, issuer, audience, now)

	revokedAt := now.Add(-10 * time.Second)
	input := ApplySessionRevokedInput{Issuer: issuer, Audience: audience, JTI: uuid.NewString(), TransactionID: uuid.NewString(), SubjectIssuer: subjectIssuer, TenantID: tenantID, Subject: subject, EventTimestamp: revokedAt, InitiatingEntity: "admin"}
	result, err := db.ApplySessionRevoked(ctx, input)
	if err != nil || !result.WatermarkMoved || result.AuditID == "" {
		t.Fatalf("apply result=%#v error=%v", result, err)
	}
	duplicate, err := db.ApplySessionRevoked(ctx, input)
	if err != nil || !duplicate.Duplicate {
		t.Fatalf("duplicate result=%#v error=%v", duplicate, err)
	}
	evaluation, _ = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, "", tenantID, subject, revokedAt, now, 180*time.Second)
	if !evaluation.Revoked {
		t.Fatal("token at watermark was accepted")
	}
	evaluation, _ = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, "https://other.test/issuer", "", tenantID, subject, revokedAt, now, 180*time.Second)
	if evaluation.Revoked {
		t.Fatal("watermark crossed OIDC issuer boundary")
	}
	evaluation, _ = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, "", tenantID, subject, now, now, 180*time.Second)
	if evaluation.Revoked {
		t.Fatal("token after watermark was rejected")
	}

	_, _ = db.pool.Exec(ctx, `UPDATE gateway_security_event_stream_state SET last_verification_at=$3
		WHERE issuer=$1 AND audience=$2`, issuer, audience, now.Add(-181*time.Second))
	evaluation, err = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, "", tenantID, subject, now, now, 180*time.Second)
	if err != nil || !evaluation.RequireIntrospection || evaluation.Mode != "introspection_fallback" {
		t.Fatalf("fallback evaluation=%#v error=%v", evaluation, err)
	}
	connectionID := uuid.NewString()
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
		t.Fatal(err)
	}
	if _, err := db.pool.Exec(ctx, `INSERT INTO gateway_security_event_connections(
		connection_id,transmitter_issuer,endpoint_url,audience,stream_id,credential_ref,
		lifecycle_status,last_error_category,idempotency_key)
		VALUES($1,$2,'https://gateway.test/api/v1/security-events/ssf',$3,$4,'test-secret-ref',
		'degraded','bootstrap_verification_stale',$5)`, connectionID, issuer, audience, streamID, uuid.NewString()); err != nil {
		t.Fatal(err)
	}
	confirmAt("recovery-verification-one", now.Add(362*time.Second))
	confirmAt("recovery-verification-two", now.Add(363*time.Second))
	var lifecycle string
	var lastError *string
	if err := db.pool.QueryRow(ctx, `SELECT lifecycle_status,last_error_category
		FROM gateway_security_event_connections WHERE connection_id=$1`, connectionID).Scan(&lifecycle, &lastError); err != nil {
		t.Fatal(err)
	}
	if lifecycle != "enabled" || lastError != nil {
		t.Fatalf("recovered connection lifecycle=%q lastError=%v", lifecycle, lastError)
	}
	replacementStreamID := uuid.NewString()
	if err := db.EnsureSecurityEventStreamState(ctx, issuer, audience, replacementStreamID); err != nil {
		t.Fatalf("rebind retained receiver state: %v", err)
	}
	var reboundStreamID, reboundMode, reboundStatus string
	var reboundVerification *time.Time
	if err := db.pool.QueryRow(ctx, `SELECT stream_id::text,mode,stream_status,last_verification_at
		FROM gateway_security_event_stream_state WHERE issuer=$1 AND audience=$2`, issuer, audience).
		Scan(&reboundStreamID, &reboundMode, &reboundStatus, &reboundVerification); err != nil {
		t.Fatal(err)
	}
	if reboundStreamID != replacementStreamID || reboundMode != "bootstrap" || reboundStatus != "unknown" || reboundVerification != nil {
		t.Fatalf("rebound state stream=%q mode=%q status=%q verification=%v",
			reboundStreamID, reboundMode, reboundStatus, reboundVerification)
	}
	var retainedReceipts, retainedWatermarks int
	if err := db.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_security_event_receipts
		WHERE issuer=$1`, issuer).Scan(&retainedReceipts); err != nil {
		t.Fatal(err)
	}
	if err := db.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_oidc_revocation_watermarks
		WHERE issuer=$1 AND tenant_id=$2 AND subject=$3`, subjectIssuer, tenantID, subject).Scan(&retainedWatermarks); err != nil {
		t.Fatal(err)
	}
	if retainedReceipts == 0 || retainedWatermarks != 1 {
		t.Fatalf("rebind discarded security history receipts=%d watermarks=%d", retainedReceipts, retainedWatermarks)
	}
}

func TestApplicationScopedSecurityEventsRevokeOnlyMatchingTenantSessions(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 application-scoped security event integration tests")
	}
	ctx := context.Background()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	transmitterIssuer := "https://auth.test.example/ssf/" + uuid.NewString()
	subjectIssuer := "https://auth.test.example/issuer/shared"
	applicationID, tenantA, tenantB, subject := uuid.NewString(), uuid.NewString(), uuid.NewString(), uuid.NewString()
	projection := func(tenantID string) *auth.User {
		result, err := db.ResolveOrProvisionOIDCUser(ctx, ResolveOrProvisionOIDCUserInput{
			Issuer: subjectIssuer, ApplicationID: applicationID, Subject: subject, Username: "shared-user",
			Roles: []string{"basic"}, TenantID: tenantID, TenantMode: "multi_tenant",
			TenantName: "Tenant " + tenantID[:8], TenantSlug: "tenant-" + tenantID[:8],
			TenantMetadataStatus: "synced", TenantMetadataVersion: "1",
			TenantMetadataUpdatedAt: time.Now().UTC(), OIDCClientID: "gateway-browser", ProvisioningEnabled: true,
		})
		if err != nil {
			t.Fatal(err)
		}
		return result.User
	}
	userA, userB := projection(tenantA), projection(tenantB)
	if _, err := db.CreateAPIKey(ctx, CreateAPIKeyInput{Name: "must survive OIDC revocation"}, userA); err != nil {
		t.Fatal(err)
	}
	createSession := func(user *auth.User, marker byte) {
		now := time.Now().UTC()
		sessionHash := sha256.Sum256([]byte(uuid.NewString()))
		if _, err := db.CreateOIDCSession(ctx, CreateOIDCSessionInput{
			SessionTokenHash: sessionHash[:], GatewayUserID: user.GatewayUserID,
			GatewayTenantID: user.GatewayTenantID, OIDCUserBindingID: user.OIDCUserBindingID,
			OIDCClientID: "gateway-browser", Issuer: subjectIssuer, ApplicationID: applicationID,
			TenantID: user.TenantID, TokenCiphertext: []byte{marker},
			AccessTokenExpiresAt: now.Add(time.Hour), LastSeenAt: now,
			IdleExpiresAt: now.Add(time.Hour), AbsoluteExpiresAt: now.Add(2 * time.Hour),
		}); err != nil {
			t.Fatal(err)
		}
	}
	createSession(userA, 0xa1)
	createSession(userB, 0xb1)

	revokedAt := time.Now().UTC().Truncate(time.Second)
	principalEvent := ApplySessionRevokedInput{
		Issuer: transmitterIssuer, Audience: "urn:easyai:ssf:receiver:" + applicationID,
		JTI: uuid.NewString(), TransactionID: uuid.NewString(), SubjectIssuer: subjectIssuer,
		ApplicationID: applicationID, SubjectType: "principal", TenantID: tenantA, Subject: subject,
		EventTimestamp: revokedAt, InitiatingEntity: "admin",
	}
	result, err := db.ApplySessionRevoked(ctx, principalEvent)
	if err != nil || result.SessionsDeleted != 1 || !result.WatermarkMoved {
		t.Fatalf("principal result=%#v error=%v", result, err)
	}
	var tenantASessions, tenantBSessions int
	countSessions := func() {
		t.Helper()
		if err := db.pool.QueryRow(ctx, `SELECT
			(SELECT count(*) FROM gateway_oidc_sessions WHERE gateway_tenant_id=$1::uuid),
			(SELECT count(*) FROM gateway_oidc_sessions WHERE gateway_tenant_id=$2::uuid)`,
			userA.GatewayTenantID, userB.GatewayTenantID,
		).Scan(&tenantASessions, &tenantBSessions); err != nil {
			t.Fatal(err)
		}
	}
	countSessions()
	if tenantASessions != 0 || tenantBSessions != 1 {
		t.Fatalf("principal revocation sessions A=%d B=%d", tenantASessions, tenantBSessions)
	}
	older := principalEvent
	older.JTI = uuid.NewString()
	older.EventTimestamp = revokedAt.Add(-time.Minute)
	result, err = db.ApplySessionRevoked(ctx, older)
	if err != nil || result.WatermarkMoved || result.SessionsDeleted != 0 {
		t.Fatalf("older result=%#v error=%v", result, err)
	}

	createSession(userA, 0xa2)
	tenantEvent := principalEvent
	tenantEvent.JTI, tenantEvent.TransactionID = uuid.NewString(), uuid.NewString()
	tenantEvent.SubjectType, tenantEvent.Subject = "tenant", tenantA
	tenantEvent.EventTimestamp = revokedAt.Add(time.Minute)
	result, err = db.ApplySessionRevoked(ctx, tenantEvent)
	if err != nil || result.SessionsDeleted != 1 || !result.WatermarkMoved {
		t.Fatalf("tenant result=%#v error=%v", result, err)
	}
	countSessions()
	if tenantASessions != 0 || tenantBSessions != 1 {
		t.Fatalf("tenant revocation sessions A=%d B=%d", tenantASessions, tenantBSessions)
	}
	keysA, err := db.ListAPIKeys(ctx, userA)
	if err != nil || len(keysA) != 1 {
		t.Fatalf("tenant revocation changed API keys: keys=%#v error=%v", keysA, err)
	}
	var accessStatus string
	if err := db.pool.QueryRow(ctx, `SELECT access_status FROM gateway_oidc_tenant_bindings
		WHERE issuer=$1 AND application_id=$2 AND external_tenant_id=$3`,
		subjectIssuer, applicationID, tenantA,
	).Scan(&accessStatus); err != nil || accessStatus != "disabled" {
		t.Fatalf("tenant A access status=%q error=%v", accessStatus, err)
	}

	streamID := uuid.NewString()
	audience := tenantEvent.Audience
	if err := db.EnsureSecurityEventStreamState(ctx, transmitterIssuer, audience, streamID); err != nil {
		t.Fatal(err)
	}
	_, _ = db.pool.Exec(ctx, `UPDATE gateway_security_event_stream_state
		SET stream_status='enabled',mode='push_healthy',last_verification_at=$3,bootstrap_until=$3
		WHERE issuer=$1 AND audience=$2`, transmitterIssuer, audience, revokedAt.Add(2*time.Minute))
	evaluationA, err := db.EvaluateOIDCSecurityEvent(
		ctx, transmitterIssuer, audience, subjectIssuer, applicationID, tenantA, subject,
		tenantEvent.EventTimestamp, revokedAt.Add(2*time.Minute), 5*time.Minute,
	)
	if err != nil || !evaluationA.Revoked {
		t.Fatalf("tenant A evaluation=%#v error=%v", evaluationA, err)
	}
	evaluationB, err := db.EvaluateOIDCSecurityEvent(
		ctx, transmitterIssuer, audience, subjectIssuer, applicationID, tenantB, subject,
		tenantEvent.EventTimestamp, revokedAt.Add(2*time.Minute), 5*time.Minute,
	)
	if err != nil || evaluationB.Revoked {
		t.Fatalf("tenant B evaluation=%#v error=%v", evaluationB, err)
	}
}