package sessionbroker

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"slices"
	"strings"
	"testing"
	"time"

	"github.com/example/remote-access-platform/backend/internal/platform/config"
	"github.com/example/remote-access-platform/backend/internal/platform/module"
	sessioncontracts "github.com/example/remote-access-platform/backend/pkg/contracts/session"
)

func TestDataPlaneTokenScopeValidation(t *testing.T) {
	now := time.Now().UTC().Truncate(time.Second)
	privateKeyPEM, publicKey := testRS256Key(t)
	service := &Service{
		cfg: module.Config{
			Auth: config.AuthConfig{
				Issuer: "rap-api-test",
			},
			DataPlane: config.DataPlaneConfig{
				TokenTTL:           time.Minute,
				TokenPrivateKeyPEM: privateKeyPEM,
				BackendGatewayURL:  "wss://backend.example.test/api/v1/gateway/ws",
			},
		},
		now: func() time.Time { return now },
	}
	session := RemoteSession{
		ID:             "session-1",
		OrganizationID: "org-1",
		ResourceID:     "resource-1",
		WorkerID:       "worker-1",
		Metadata:       mustJSON(t, map[string]any{"policy": map[string]any{"clipboard_mode": "bidirectional", "file_transfer_mode": "client_to_server"}}),
	}
	attachment := SessionAttachment{
		ID:     "attachment-1",
		UserID: "user-1",
	}

	offer, err := service.buildDataPlaneOffer(session, attachment)
	if err != nil {
		t.Fatalf("buildDataPlaneOffer returned error: %v", err)
	}
	if offer == nil {
		t.Fatal("expected data-plane offer")
	}

	claims, err := parseDataPlaneToken(offer.Token, publicKey)
	if err != nil {
		t.Fatalf("parseDataPlaneToken returned error: %v", err)
	}
	assertEqual(t, claims.SessionID, session.ID, "session_id")
	assertEqual(t, claims.AttachmentID, attachment.ID, "attachment_id")
	assertEqual(t, claims.UserID, attachment.UserID, "user_id")
	assertEqual(t, claims.OrganizationID, session.OrganizationID, "organization_id")
	assertEqual(t, claims.WorkerID, session.WorkerID, "worker_id")
	assertEqual(t, claims.ResourceID, session.ResourceID, "resource_id")
	if claims.ID == "" {
		t.Fatal("expected jti")
	}
	if claims.ExpiresAt == nil || !claims.ExpiresAt.Time.Equal(now.Add(time.Minute)) {
		t.Fatalf("unexpected expires_at: %v", claims.ExpiresAt)
	}
	if !claims.ExpiresAtValue.Equal(now.Add(time.Minute)) {
		t.Fatalf("unexpected expires_at claim value: %v", claims.ExpiresAtValue)
	}
	for _, channel := range []string{
		sessioncontracts.DataPlaneChannelControl,
		sessioncontracts.DataPlaneChannelInput,
		sessioncontracts.DataPlaneChannelRender,
		sessioncontracts.DataPlaneChannelTelemetry,
		sessioncontracts.DataPlaneChannelClipboard,
		sessioncontracts.DataPlaneChannelFileUpload,
	} {
		if !slices.Contains(claims.AllowedChannels, channel) {
			t.Fatalf("expected allowed channel %q in %v", channel, claims.AllowedChannels)
		}
	}
}

func TestDataPlaneOfferResponseShapeCompatibility(t *testing.T) {
	now := time.Now().UTC().Truncate(time.Second)
	privateKeyPEM, _ := testRS256Key(t)
	service := &Service{
		cfg: module.Config{
			Auth: config.AuthConfig{Issuer: "rap-api-test"},
			DataPlane: config.DataPlaneConfig{
				TokenTTL:                   time.Minute,
				TokenPrivateKeyPEM:         privateKeyPEM,
				BackendGatewayURL:          "wss://backend.example.test/api/v1/gateway/ws",
				DirectWorkerWSSURLTemplate: "wss://{worker_id}.worker.example.test/rap/v1/data-plane",
				DirectWorkerJSONRuntime:    true,
				DirectWorkerTLSTrustMode:   "smoke_insecure",
			},
		},
		now: func() time.Time { return now },
	}
	result := &SessionControlResult{
		Session: RemoteSession{
			ID:             "session-1",
			OrganizationID: "org-1",
			ResourceID:     "resource-1",
			WorkerID:       "worker-1",
			Metadata:       mustJSON(t, map[string]any{"policy": map[string]any{"clipboard_mode": "disabled", "file_transfer_mode": "disabled"}}),
		},
		Attachment: &SessionAttachment{ID: "attachment-1", UserID: "user-1"},
		AttachToken: &sessioncontracts.AttachTokenClaims{
			Token:        "existing-attach-token",
			SessionID:    "session-1",
			AttachmentID: "attachment-1",
			UserID:       "user-1",
			WorkerID:     "worker-1",
			ExpiresAt:    now.Add(2 * time.Minute),
		},
	}

	if err := service.attachDataPlaneOffer(result); err != nil {
		t.Fatalf("attachDataPlaneOffer returned error: %v", err)
	}
	payload, err := json.Marshal(result)
	if err != nil {
		t.Fatalf("marshal response: %v", err)
	}
	var decoded map[string]any
	if err := json.Unmarshal(payload, &decoded); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if decoded["session"] == nil || decoded["attachment"] == nil || decoded["attach_token"] == nil {
		t.Fatalf("response lost existing fields: %s", payload)
	}
	if decoded["data_plane"] == nil || decoded["gateway_url"] == nil {
		t.Fatalf("response missing data-plane fields: %s", payload)
	}
	if result.DataPlane == nil {
		t.Fatal("expected data-plane offer")
	}
	if result.DataPlane.Preferred != sessioncontracts.DataPlaneCandidateDirectWorkerWSS {
		t.Fatalf("unexpected preferred candidate: %s", result.DataPlane.Preferred)
	}
	if len(result.DataPlane.Candidates) != 2 {
		t.Fatalf("expected direct and fallback candidates, got %d", len(result.DataPlane.Candidates))
	}
	if result.DataPlane.Candidates[0].URL != "wss://worker-1.worker.example.test/rap/v1/data-plane" {
		t.Fatalf("unexpected direct candidate URL: %s", result.DataPlane.Candidates[0].URL)
	}
	if result.DataPlane.Candidates[0].Metadata["runtime_transport"] != "json_v1" {
		t.Fatalf("direct candidate is missing json_v1 runtime metadata: %#v", result.DataPlane.Candidates[0].Metadata)
	}
	if result.DataPlane.Candidates[0].Metadata["traffic_ready"] != true {
		t.Fatalf("direct candidate is missing traffic_ready metadata: %#v", result.DataPlane.Candidates[0].Metadata)
	}
	if result.DataPlane.Candidates[0].Metadata["smoke_only"] != true {
		t.Fatalf("direct candidate should be marked smoke-only by default: %#v", result.DataPlane.Candidates[0].Metadata)
	}
	if result.DataPlane.Candidates[0].Metadata["production_trusted"] != false {
		t.Fatalf("smoke direct candidate must not be production-trusted: %#v", result.DataPlane.Candidates[0].Metadata)
	}
	if !strings.Contains(result.DataPlane.Candidates[1].URL, "/api/v1/gateway/ws") {
		t.Fatalf("unexpected backend candidate URL: %s", result.DataPlane.Candidates[1].URL)
	}
}

func TestDataPlaneDirectCandidateMetadataRequiresRuntimeFlag(t *testing.T) {
	service := &Service{
		cfg: module.Config{
			DataPlane: config.DataPlaneConfig{
				BackendGatewayURL:          "wss://backend.example.test/api/v1/gateway/ws",
				DirectWorkerWSSURLTemplate: "wss://{worker_id}.worker.example.test/rap/v1/data-plane",
				DirectWorkerTLSTrustMode:   "smoke_insecure",
			},
		},
	}
	candidates := service.buildDataPlaneCandidates(RemoteSession{WorkerID: "worker-1"})
	if len(candidates) != 2 {
		t.Fatalf("expected direct and fallback candidates, got %d", len(candidates))
	}
	if candidates[0].Metadata != nil {
		t.Fatalf("direct candidate must not advertise json_v1 before runtime flag is enabled: %#v", candidates[0].Metadata)
	}
}

func TestDataPlaneDirectCandidateAdvertisesBinaryRenderOnlyWhenEnabled(t *testing.T) {
	service := &Service{
		cfg: module.Config{
			DataPlane: config.DataPlaneConfig{
				BackendGatewayURL:          "wss://backend.example.test/api/v1/gateway/ws",
				DirectWorkerWSSURLTemplate: "wss://{worker_id}.worker.example.test/rap/v1/data-plane",
				DirectWorkerJSONRuntime:    true,
				DirectWorkerBinaryRender:   true,
				DirectWorkerTLSTrustMode:   "platform_ca",
				DirectWorkerTLSCARef:       "rap-platform-ca:v1",
			},
		},
	}
	candidates := service.buildDataPlaneCandidates(RemoteSession{WorkerID: "worker-1"})
	if len(candidates) != 2 {
		t.Fatalf("expected direct and fallback candidates, got %d", len(candidates))
	}
	if candidates[0].Metadata["render_transport"] != "binary_v1" {
		t.Fatalf("direct candidate is missing binary render metadata: %#v", candidates[0].Metadata)
	}
	if candidates[0].Metadata["binary_render"] != true {
		t.Fatalf("direct candidate is missing binary_render metadata: %#v", candidates[0].Metadata)
	}
	if candidates[0].Metadata["default_color_mode"] != "full_color" {
		t.Fatalf("direct candidate is missing default_color_mode metadata: %#v", candidates[0].Metadata)
	}
	if candidates[0].Metadata["production_trusted"] != true || candidates[0].Metadata["tls_trust_mode"] != "platform_ca" {
		t.Fatalf("direct candidate is missing production trust metadata: %#v", candidates[0].Metadata)
	}
	if candidates[0].Metadata["tls_ca_ref"] != "rap-platform-ca:v1" {
		t.Fatalf("direct candidate is missing tls_ca_ref metadata: %#v", candidates[0].Metadata)
	}
	modes, ok := candidates[0].Metadata["supported_color_modes"].([]string)
	if !ok || !slices.Contains(modes, "full_color") || !slices.Contains(modes, "grayscale") {
		t.Fatalf("direct candidate is missing supported_color_modes metadata: %#v", candidates[0].Metadata)
	}
}

func TestDataPlaneDirectCandidateOmittedInProductionWhenSmokeOnly(t *testing.T) {
	service := &Service{
		cfg: module.Config{
			App: config.AppConfig{Env: "production"},
			DataPlane: config.DataPlaneConfig{
				BackendGatewayURL:          "wss://backend.example.test/api/v1/gateway/ws",
				DirectWorkerWSSURLTemplate: "wss://{worker_id}.worker.example.test/rap/v1/data-plane",
				DirectWorkerJSONRuntime:    true,
				DirectWorkerTLSTrustMode:   "smoke_insecure",
			},
		},
	}
	candidates := service.buildDataPlaneCandidates(RemoteSession{WorkerID: "worker-1"})
	if len(candidates) != 1 {
		t.Fatalf("expected fallback-only candidates in production with smoke TLS, got %d", len(candidates))
	}
	if candidates[0].Type != sessioncontracts.DataPlaneCandidateBackendGateway {
		t.Fatalf("production must not advertise smoke-only direct candidate: %#v", candidates)
	}
}

func TestDataPlaneDirectCandidateAdvertisedInProductionWhenTrusted(t *testing.T) {
	service := &Service{
		cfg: module.Config{
			App: config.AppConfig{Env: "production"},
			DataPlane: config.DataPlaneConfig{
				BackendGatewayURL:          "wss://backend.example.test/api/v1/gateway/ws",
				DirectWorkerWSSURLTemplate: "wss://{worker_id}.worker.example.test/rap/v1/data-plane",
				DirectWorkerJSONRuntime:    true,
				DirectWorkerTLSTrustMode:   "public_ca",
			},
		},
	}
	candidates := service.buildDataPlaneCandidates(RemoteSession{WorkerID: "worker-1"})
	if len(candidates) != 2 {
		t.Fatalf("expected trusted direct and fallback candidates, got %d", len(candidates))
	}
	if candidates[0].Metadata["production_trusted"] != true || candidates[0].Metadata["tls_trust_mode"] != "public_ca" {
		t.Fatalf("trusted production direct candidate metadata mismatch: %#v", candidates[0].Metadata)
	}
}

func TestDataPlaneCandidatesFallbackOnlyWhenDirectTemplateMissing(t *testing.T) {
	service := &Service{
		cfg: module.Config{
			DataPlane: config.DataPlaneConfig{
				BackendGatewayURL: "wss://backend.example.test/api/v1/gateway/ws",
			},
		},
	}
	candidates := service.buildDataPlaneCandidates(RemoteSession{WorkerID: "worker-1"})
	if len(candidates) != 1 {
		t.Fatalf("expected fallback-only candidate list, got %d", len(candidates))
	}
	if candidates[0].Type != sessioncontracts.DataPlaneCandidateBackendGateway {
		t.Fatalf("unexpected candidate type: %s", candidates[0].Type)
	}
}

func TestDataPlaneAllowedChannelsRespectRuntimePolicy(t *testing.T) {
	cases := []struct {
		name     string
		policy   map[string]any
		expected []string
		blocked  []string
	}{
		{
			name:     "disabled policies expose only control input render telemetry",
			policy:   map[string]any{"clipboard_mode": "disabled", "file_transfer_mode": "disabled"},
			expected: []string{sessioncontracts.DataPlaneChannelControl, sessioncontracts.DataPlaneChannelInput, sessioncontracts.DataPlaneChannelRender, sessioncontracts.DataPlaneChannelTelemetry},
			blocked:  []string{sessioncontracts.DataPlaneChannelClipboard, sessioncontracts.DataPlaneChannelFileUpload},
		},
		{
			name:     "clipboard policy adds clipboard channel",
			policy:   map[string]any{"clipboard_mode": "server_to_client", "file_transfer_mode": "disabled"},
			expected: []string{sessioncontracts.DataPlaneChannelClipboard},
			blocked:  []string{sessioncontracts.DataPlaneChannelFileUpload},
		},
		{
			name:     "client upload policy adds file upload channel",
			policy:   map[string]any{"clipboard_mode": "disabled", "file_transfer_mode": "client_to_server"},
			expected: []string{sessioncontracts.DataPlaneChannelFileUpload},
			blocked:  []string{sessioncontracts.DataPlaneChannelClipboard},
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			session := RemoteSession{Metadata: mustJSON(t, map[string]any{"policy": tc.policy})}
			channels := dataPlaneAllowedChannelsFromSession(session)
			for _, channel := range tc.expected {
				if !slices.Contains(channels, channel) {
					t.Fatalf("expected channel %q in %v", channel, channels)
				}
			}
			for _, channel := range tc.blocked {
				if slices.Contains(channels, channel) {
					t.Fatalf("did not expect channel %q in %v", channel, channels)
				}
			}
		})
	}
}

func mustJSON(t *testing.T, value any) []byte {
	t.Helper()
	payload, err := json.Marshal(value)
	if err != nil {
		t.Fatalf("marshal test metadata: %v", err)
	}
	return payload
}

func testRS256Key(t *testing.T) (string, *rsa.PublicKey) {
	t.Helper()
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate RSA key: %v", err)
	}
	encoded := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	})
	return string(encoded), &privateKey.PublicKey
}

func assertEqual(t *testing.T, got, want, name string) {
	t.Helper()
	if got != want {
		t.Fatalf("unexpected %s: got %q want %q", name, got, want)
	}
}