package dht

import (
	"context"
	"testing"
	"time"

	"chorus/pkg/config"
)

// TestDHTSecurityPolicyEnforcement tests security policy enforcement in DHT operations
func TestDHTSecurityPolicyEnforcement(t *testing.T) {
	ctx := context.Background()
	
	testCases := []struct {
		name            string
		currentRole     string
		operation       string
		ucxlAddress     string
		contentType     string
		expectSuccess   bool
		expectedError   string
	}{
		// Store operation tests
		{
			name:          "admin_can_store_all_content",
			currentRole:   "admin",
			operation:     "store",
			ucxlAddress:   "agent1:admin:system:security_audit",
			contentType:   "decision",
			expectSuccess: true,
		},
		{
			name:          "backend_developer_can_store_backend_content",
			currentRole:   "backend_developer",
			operation:     "store", 
			ucxlAddress:   "agent1:backend_developer:api:endpoint_design",
			contentType:   "suggestion",
			expectSuccess: true,
		},
		{
			name:            "readonly_role_cannot_store",
			currentRole:     "readonly_user",
			operation:       "store",
			ucxlAddress:     "agent1:readonly_user:project:observation",
			contentType:     "suggestion",
			expectSuccess:   false,
			expectedError:   "read-only authority",
		},
		{
			name:            "unknown_role_cannot_store",
			currentRole:     "invalid_role",
			operation:       "store",
			ucxlAddress:     "agent1:invalid_role:project:task",
			contentType:     "decision",
			expectSuccess:   false,
			expectedError:   "unknown creator role",
		},
		
		// Retrieve operation tests
		{
			name:          "any_valid_role_can_retrieve",
			currentRole:   "qa_engineer",
			operation:     "retrieve",
			ucxlAddress:   "agent1:backend_developer:api:test_data",
			expectSuccess: true,
		},
		{
			name:            "unknown_role_cannot_retrieve",
			currentRole:     "nonexistent_role",
			operation:       "retrieve",
			ucxlAddress:     "agent1:backend_developer:api:test_data",
			expectSuccess:   false,
			expectedError:   "unknown current role",
		},
		
		// Announce operation tests
		{
			name:          "coordination_role_can_announce",
			currentRole:   "senior_software_architect",
			operation:     "announce",
			ucxlAddress:   "agent1:senior_software_architect:architecture:blueprint",
			expectSuccess: true,
		},
		{
			name:          "decision_role_can_announce",
			currentRole:   "security_expert",
			operation:     "announce",
			ucxlAddress:   "agent1:security_expert:security:policy",
			expectSuccess: true,
		},
		{
			name:            "suggestion_role_cannot_announce",
			currentRole:     "suggestion_only_role",
			operation:       "announce",
			ucxlAddress:     "agent1:suggestion_only_role:project:idea",
			expectSuccess:   false,
			expectedError:   "lacks authority",
		},
		{
			name:            "readonly_role_cannot_announce",
			currentRole:     "readonly_user",
			operation:       "announce",
			ucxlAddress:     "agent1:readonly_user:project:observation",
			expectSuccess:   false,
			expectedError:   "lacks authority",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Create test configuration
			cfg := &config.Config{
				Agent: config.AgentConfig{
					ID:   "test-agent",
					Role: tc.currentRole,
				},
				Security: config.SecurityConfig{
					KeyRotationDays: 90,
					AuditLogging:    true,
					AuditPath:       "/tmp/test-security-audit.log",
				},
			}

			// Create mock encrypted storage
			eds := createMockEncryptedStorage(ctx, cfg)

			var err error
			switch tc.operation {
			case "store":
				err = eds.checkStoreAccessPolicy(tc.currentRole, tc.ucxlAddress, tc.contentType)
			case "retrieve":
				err = eds.checkRetrieveAccessPolicy(tc.currentRole, tc.ucxlAddress)
			case "announce":
				err = eds.checkAnnounceAccessPolicy(tc.currentRole, tc.ucxlAddress)
			}

			if tc.expectSuccess {
				if err != nil {
					t.Errorf("Expected %s operation to succeed for role %s, but got error: %v", 
						tc.operation, tc.currentRole, err)
				}
			} else {
				if err == nil {
					t.Errorf("Expected %s operation to fail for role %s, but it succeeded", 
						tc.operation, tc.currentRole)
				}
				if tc.expectedError != "" && !containsSubstring(err.Error(), tc.expectedError) {
					t.Errorf("Expected error to contain '%s', got '%s'", tc.expectedError, err.Error())
				}
			}
		})
	}
}

// TestDHTAuditLogging tests comprehensive audit logging for DHT operations
func TestDHTAuditLogging(t *testing.T) {
	ctx := context.Background()
	
	testCases := []struct {
		name         string
		operation    string
		role         string
		ucxlAddress  string
		success      bool
		errorMsg     string
		expectAudit  bool
	}{
		{
			name:        "successful_store_operation",
			operation:   "store",
			role:        "backend_developer", 
			ucxlAddress: "agent1:backend_developer:api:user_service",
			success:     true,
			expectAudit: true,
		},
		{
			name:        "failed_store_operation",
			operation:   "store",
			role:        "readonly_user",
			ucxlAddress: "agent1:readonly_user:project:readonly_attempt",
			success:     false,
			errorMsg:    "read-only authority",
			expectAudit: true,
		},
		{
			name:        "successful_retrieve_operation",
			operation:   "retrieve",
			role:        "frontend_developer",
			ucxlAddress: "agent1:backend_developer:api:user_data",
			success:     true,
			expectAudit: true,
		},
		{
			name:        "successful_announce_operation",
			operation:   "announce",
			role:        "senior_software_architect",
			ucxlAddress: "agent1:senior_software_architect:architecture:system_design",
			success:     true,
			expectAudit: true,
		},
		{
			name:        "audit_disabled_no_logging",
			operation:   "store",
			role:        "backend_developer",
			ucxlAddress: "agent1:backend_developer:api:no_audit",
			success:     true,
			expectAudit: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Create configuration with audit logging
			cfg := &config.Config{
				Agent: config.AgentConfig{
					ID:   "test-agent",
					Role: tc.role,
				},
				Security: config.SecurityConfig{
					KeyRotationDays: 90,
					AuditLogging:    tc.expectAudit,
					AuditPath:       "/tmp/test-dht-audit.log",
				},
			}

			// Create mock encrypted storage
			eds := createMockEncryptedStorage(ctx, cfg)
			
			// Capture audit output
			auditCaptured := false

			// Simulate audit operation
			switch tc.operation {
			case "store":
				// Mock the audit function call
				if tc.expectAudit && cfg.Security.AuditLogging {
					eds.auditStoreOperation(tc.ucxlAddress, tc.role, "test-content", 1024, tc.success, tc.errorMsg)
					auditCaptured = true
				}
			case "retrieve":
				if tc.expectAudit && cfg.Security.AuditLogging {
					eds.auditRetrieveOperation(tc.ucxlAddress, tc.role, tc.success, tc.errorMsg)
					auditCaptured = true
				}
			case "announce":
				if tc.expectAudit && cfg.Security.AuditLogging {
					eds.auditAnnounceOperation(tc.ucxlAddress, tc.role, tc.success, tc.errorMsg)
					auditCaptured = true
				}
			}

			// Verify audit logging behavior
			if tc.expectAudit && !auditCaptured {
				t.Errorf("Expected audit logging for %s operation but none was captured", tc.operation)
			}
			if !tc.expectAudit && auditCaptured {
				t.Errorf("Expected no audit logging for %s operation but audit was captured", tc.operation)
			}
		})
	}
}

// TestSecurityConfigIntegration tests integration with SecurityConfig
func TestSecurityConfigIntegration(t *testing.T) {
	ctx := context.Background()
	
	testConfigs := []struct {
		name            string
		auditLogging    bool
		auditPath       string
		expectAuditWork bool
	}{
		{
			name:            "audit_enabled_with_path",
			auditLogging:    true,
			auditPath:       "/tmp/test-audit-enabled.log",
			expectAuditWork: true,
		},
		{
			name:            "audit_disabled",
			auditLogging:    false,
			auditPath:       "/tmp/test-audit-disabled.log",
			expectAuditWork: false,
		},
		{
			name:            "audit_enabled_no_path",
			auditLogging:    true,
			auditPath:       "",
			expectAuditWork: false,
		},
	}

	for _, tc := range testConfigs {
		t.Run(tc.name, func(t *testing.T) {
			cfg := &config.Config{
				Agent: config.AgentConfig{
					ID:   "test-agent",
					Role: "backend_developer",
				},
				Security: config.SecurityConfig{
					KeyRotationDays: 90,
					AuditLogging:    tc.auditLogging,
					AuditPath:       tc.auditPath,
				},
			}

			eds := createMockEncryptedStorage(ctx, cfg)

			// Test audit function behavior with different configurations
			auditWorked := func() bool {
				if !cfg.Security.AuditLogging || cfg.Security.AuditPath == "" {
					return false
				}
				return true
			}()

			if auditWorked != tc.expectAuditWork {
				t.Errorf("Expected audit to work: %v, but got: %v", tc.expectAuditWork, auditWorked)
			}
		})
	}
}

// TestRoleAuthorityHierarchy tests role authority hierarchy enforcement
func TestRoleAuthorityHierarchy(t *testing.T) {
	ctx := context.Background()
	
	// Test role authority levels for different operations
	authorityTests := []struct {
		role            string
		authorityLevel  config.AuthorityLevel
		canStore        bool
		canRetrieve     bool  
		canAnnounce     bool
	}{
		{
			role:            "admin",
			authorityLevel:  config.AuthorityMaster,
			canStore:        true,
			canRetrieve:     true,
			canAnnounce:     true,
		},
		{
			role:            "senior_software_architect",
			authorityLevel:  config.AuthorityDecision,
			canStore:        true,
			canRetrieve:     true,
			canAnnounce:     true,
		},
		{
			role:            "security_expert",
			authorityLevel:  config.AuthorityCoordination,
			canStore:        true,
			canRetrieve:     true,
			canAnnounce:     true,
		},
		{
			role:            "backend_developer",
			authorityLevel:  config.AuthoritySuggestion,
			canStore:        true,
			canRetrieve:     true,
			canAnnounce:     false,
		},
	}

	for _, tt := range authorityTests {
		t.Run(tt.role+"_authority_test", func(t *testing.T) {
			cfg := &config.Config{
				Agent: config.AgentConfig{
					ID:   "test-agent", 
					Role: tt.role,
				},
				Security: config.SecurityConfig{
					KeyRotationDays: 90,
					AuditLogging:    true,
					AuditPath:       "/tmp/test-authority.log",
				},
			}

			eds := createMockEncryptedStorage(ctx, cfg)

			// Test store permission
			storeErr := eds.checkStoreAccessPolicy(tt.role, "test:address", "content")
			if tt.canStore && storeErr != nil {
				t.Errorf("Role %s should be able to store but got error: %v", tt.role, storeErr)
			}
			if !tt.canStore && storeErr == nil {
				t.Errorf("Role %s should not be able to store but operation succeeded", tt.role)
			}

			// Test retrieve permission
			retrieveErr := eds.checkRetrieveAccessPolicy(tt.role, "test:address")
			if tt.canRetrieve && retrieveErr != nil {
				t.Errorf("Role %s should be able to retrieve but got error: %v", tt.role, retrieveErr)
			}
			if !tt.canRetrieve && retrieveErr == nil {
				t.Errorf("Role %s should not be able to retrieve but operation succeeded", tt.role)
			}

			// Test announce permission
			announceErr := eds.checkAnnounceAccessPolicy(tt.role, "test:address")
			if tt.canAnnounce && announceErr != nil {
				t.Errorf("Role %s should be able to announce but got error: %v", tt.role, announceErr)
			}
			if !tt.canAnnounce && announceErr == nil {
				t.Errorf("Role %s should not be able to announce but operation succeeded", tt.role)
			}
		})
	}
}

// TestSecurityMetrics tests security-related metrics
func TestSecurityMetrics(t *testing.T) {
	ctx := context.Background()
	
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID:   "test-agent",
			Role: "backend_developer",
		},
		Security: config.SecurityConfig{
			KeyRotationDays: 90,
			AuditLogging:    true,
			AuditPath:       "/tmp/test-metrics.log",
		},
	}

	eds := createMockEncryptedStorage(ctx, cfg)

	// Simulate some operations to generate metrics
	for i := 0; i < 5; i++ {
		eds.metrics.StoredItems++
		eds.metrics.RetrievedItems++
		eds.metrics.EncryptionOps++
		eds.metrics.DecryptionOps++
	}

	metrics := eds.GetMetrics()

	expectedMetrics := map[string]int64{
		"stored_items":    5,
		"retrieved_items": 5,
		"encryption_ops":  5,
		"decryption_ops":  5,
	}

	for metricName, expectedValue := range expectedMetrics {
		if actualValue, ok := metrics[metricName]; !ok {
			t.Errorf("Expected metric %s to be present in metrics", metricName)
		} else if actualValue != expectedValue {
			t.Errorf("Expected %s to be %d, got %v", metricName, expectedValue, actualValue)
		}
	}
}

// Helper functions

func createMockEncryptedStorage(ctx context.Context, cfg *config.Config) *EncryptedDHTStorage {
	return &EncryptedDHTStorage{
		ctx:     ctx,
		config:  cfg,
		nodeID:  "test-node-id",
		cache:   make(map[string]*CachedEntry),
		metrics: &StorageMetrics{
			LastUpdate: time.Now(),
		},
	}
}

func containsSubstring(str, substr string) bool {
	if len(substr) == 0 {
		return true
	}
	if len(str) < len(substr) {
		return false
	}
	for i := 0; i <= len(str)-len(substr); i++ {
		if str[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}

// Benchmarks for security performance

func BenchmarkSecurityPolicyChecks(b *testing.B) {
	ctx := context.Background()
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID:   "bench-agent",
			Role: "backend_developer",
		},
		Security: config.SecurityConfig{
			KeyRotationDays: 90,
			AuditLogging:    true,
			AuditPath:       "/tmp/bench-security.log",
		},
	}

	eds := createMockEncryptedStorage(ctx, cfg)

	b.ResetTimer()

	b.Run("store_policy_check", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			eds.checkStoreAccessPolicy("backend_developer", "test:address", "content")
		}
	})

	b.Run("retrieve_policy_check", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			eds.checkRetrieveAccessPolicy("backend_developer", "test:address")
		}
	})

	b.Run("announce_policy_check", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			eds.checkAnnounceAccessPolicy("senior_software_architect", "test:address")
		}
	})
}

func BenchmarkAuditOperations(b *testing.B) {
	ctx := context.Background()
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID:   "bench-agent",
			Role: "backend_developer",
		},
		Security: config.SecurityConfig{
			KeyRotationDays: 90,
			AuditLogging:    true,
			AuditPath:       "/tmp/bench-audit.log",
		},
	}

	eds := createMockEncryptedStorage(ctx, cfg)

	b.ResetTimer()

	b.Run("store_audit", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			eds.auditStoreOperation("test:address", "backend_developer", "content", 1024, true, "")
		}
	})

	b.Run("retrieve_audit", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			eds.auditRetrieveOperation("test:address", "backend_developer", true, "")
		}
	})

	b.Run("announce_audit", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			eds.auditAnnounceOperation("test:address", "backend_developer", true, "")
		}
	})
}