package crypto

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"testing"
	"time"

	"chorus/pkg/config"
)

// TestSecurityConfig tests SecurityConfig enforcement
func TestSecurityConfig(t *testing.T) {
	// Create temporary audit log file
	tmpDir, err := ioutil.TempDir("", "chorus_security_test")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	// Test cases for security configuration
	testCases := []struct {
		name                string
		keyRotationDays     int
		auditLogging        bool
		expectWarnings      int
		expectRotationJobs  bool
	}{
		{
			name:               "audit_logging_disabled",
			keyRotationDays:    90,
			auditLogging:       false,
			expectWarnings:     1, // Warning for disabled audit logging
			expectRotationJobs: true,
		},
		{
			name:               "key_rotation_disabled", 
			keyRotationDays:    0,
			auditLogging:       true,
			expectWarnings:     1, // Warning for disabled key rotation
			expectRotationJobs: false,
		},
		{
			name:               "security_fully_enabled",
			keyRotationDays:    30,
			auditLogging:       true,
			expectWarnings:     0,
			expectRotationJobs: true,
		},
		{
			name:               "both_security_features_disabled",
			keyRotationDays:    0,
			auditLogging:       false,
			expectWarnings:     2, // Warnings for both disabled features
			expectRotationJobs: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Create test configuration
			cfg := &config.Config{
				Agent: config.AgentConfig{
					ID: "test-agent",
				},
				Security: config.SecurityConfig{
					KeyRotationDays: tc.keyRotationDays,
					AuditLogging:    tc.auditLogging,
					AuditPath:       fmt.Sprintf("%s/audit-%s.log", tmpDir, tc.name),
				},
			}

			// Create mock audit logger
			mockLogger := &MockAuditLogger{events: make([]*SecurityEvent, 0)}
			
			// Create mock key store
			mockKeyStore := &MockKeyStore{
				keys: make(map[string]*SecureKeyData),
			}

			// Create key manager
			km, err := NewKeyManager(cfg, mockKeyStore, mockLogger)
			if err != nil {
				t.Fatalf("Failed to create key manager: %v", err)
			}
			defer func() {
				if km.rotationScheduler.running {
					km.rotationScheduler.Stop()
				}
			}()

			// Give the key manager time to initialize
			time.Sleep(100 * time.Millisecond)

			// Check audit logger for expected warnings
			securityWarnings := 0
			for _, event := range mockLogger.events {
				if event.EventType == "security_warning" {
					securityWarnings++
				}
			}

			if securityWarnings != tc.expectWarnings {
				t.Errorf("Expected %d security warnings, got %d", tc.expectWarnings, securityWarnings)
			}

			// Check if rotation scheduler is running
			isRunning := km.rotationScheduler.running
			if tc.expectRotationJobs && !isRunning {
				t.Errorf("Expected rotation scheduler to be running")
			} else if !tc.expectRotationJobs && isRunning {
				t.Errorf("Expected rotation scheduler to not be running")
			}

			// Test key rotation monitoring
			if tc.keyRotationDays > 0 {
				testKeyRotationMonitoring(t, km, mockKeyStore, mockLogger)
			}
		})
	}
}

// testKeyRotationMonitoring tests the key rotation monitoring functionality
func testKeyRotationMonitoring(t *testing.T, km *KeyManager, keyStore *MockKeyStore, mockLogger *MockAuditLogger) {
	// Create an old key that should trigger rotation warning
	oldKey := &SecureKeyData{
		KeyID:     "old-test-key",
		KeyType:   "age-x25519",
		CreatedAt: time.Now().Add(-100 * 24 * time.Hour), // 100 days old
		Status:    KeyStatusActive,
	}
	keyStore.keys[oldKey.KeyID] = oldKey

	// Create metadata for the old key
	oldKeyMeta := &KeyMetadata{
		KeyID:     "old-test-key",
		KeyType:   "age-x25519",
		RoleID:    "test-role",
		CreatedAt: time.Now().Add(-100 * 24 * time.Hour),
		Status:    KeyStatusActive,
	}
	keyStore.metadata = append(keyStore.metadata, oldKeyMeta)

	// Run key rotation check
	km.checkKeysForRotation()

	// Give time for async operations
	time.Sleep(100 * time.Millisecond)

	// Check if rotation warning was logged
	rotationWarnings := 0
	for _, event := range mockLogger.keyRotationEvents {
		if event.Reason == "key_rotation_overdue" {
			rotationWarnings++
		}
	}

	if rotationWarnings == 0 {
		t.Errorf("Expected at least one key rotation warning for overdue key")
	}
}

// TestDHTSecurityIntegration tests DHT security integration
func TestDHTSecurityIntegration(t *testing.T) {
	// Create test configuration
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID:   "test-agent",
			Role: "backend_developer",
		},
		Security: config.SecurityConfig{
			KeyRotationDays: 90,
			AuditLogging:    true,
			AuditPath:       "/tmp/test-audit.log",
		},
	}

	// Create mock DHT storage (simplified for testing)
	ctx := context.Background()
	
	// Test role-based access policies
	testCases := []struct {
		name         string
		currentRole  string
		operation    string
		shouldAllow  bool
		expectedError string
	}{
		{
			name:        "admin_can_store",
			currentRole: "admin",
			operation:   "store",
			shouldAllow: true,
		},
		{
			name:        "backend_developer_can_store",
			currentRole: "backend_developer", 
			operation:   "store",
			shouldAllow: true,
		},
		{
			name:          "readonly_cannot_store",
			currentRole:   "readonly_user",
			operation:     "store",
			shouldAllow:   false,
			expectedError: "read-only authority",
		},
		{
			name:        "all_roles_can_retrieve",
			currentRole: "qa_engineer",
			operation:   "retrieve",
			shouldAllow: true,
		},
		{
			name:          "suggestion_role_cannot_announce",
			currentRole:   "suggestion_role",
			operation:     "announce",
			shouldAllow:   false,
			expectedError: "lacks authority",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Set role in config
			cfg.Agent.Role = tc.currentRole

			// Test the specific access policy check
			var err error
			switch tc.operation {
			case "store":
				err = checkStoreAccessPolicyTest(tc.currentRole)
			case "retrieve":
				err = checkRetrieveAccessPolicyTest(tc.currentRole)
			case "announce":
				err = checkAnnounceAccessPolicyTest(tc.currentRole)
			}

			if tc.shouldAllow {
				if err != nil {
					t.Errorf("Expected operation to be allowed but got error: %v", err)
				}
			} else {
				if err == nil {
					t.Errorf("Expected operation to be denied but it was allowed")
				} else if tc.expectedError != "" && err.Error() != tc.expectedError {
					// Check if error message contains expected substring
					if len(tc.expectedError) > 0 && !containsSubstring(err.Error(), tc.expectedError) {
						t.Errorf("Expected error to contain '%s', got '%s'", tc.expectedError, err.Error())
					}
				}
			}
		})
	}
}

// TestAuditLogging tests comprehensive audit logging
func TestAuditLogging(t *testing.T) {
	tmpDir, err := ioutil.TempDir("", "chorus_audit_test")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	// Test audit logging for different operations
	testOperations := []struct {
		operation    string
		ucxlAddress  string
		role         string
		success      bool
		errorMsg     string
	}{
		{"store", "agent1:backend_developer:project1:task1", "backend_developer", true, ""},
		{"store", "agent2:invalid_role:project2:task2", "invalid_role", false, "unknown role"},
		{"retrieve", "agent1:backend_developer:project1:task1", "frontend_developer", true, ""},
		{"announce", "agent1:backend_developer:project1:task1", "senior_software_architect", true, ""},
		{"announce", "agent2:readonly:project2:task2", "readonly_user", false, "lacks authority"},
	}

	for _, op := range testOperations {
		t.Run(fmt.Sprintf("%s_%s_%v", op.operation, op.role, op.success), func(t *testing.T) {
			// Create configuration with audit logging enabled
			cfg := &config.Config{
				Agent: config.AgentConfig{
					ID:   "test-agent",
					Role: op.role,
				},
				Security: config.SecurityConfig{
					KeyRotationDays: 90,
					AuditLogging:    true,
					AuditPath:       fmt.Sprintf("%s/audit-%s.log", tmpDir, op.operation),
				},
			}

			// Simulate audit logging for the operation
			auditResult := simulateAuditOperation(cfg, op.operation, op.ucxlAddress, op.role, op.success, op.errorMsg)
			
			// Validate audit log entry
			if auditResult == nil {
				t.Errorf("Expected audit log entry but got nil")
				return
			}

			if auditResult["operation"] != op.operation {
				t.Errorf("Expected operation '%s', got '%s'", op.operation, auditResult["operation"])
			}

			if auditResult["role"] != op.role {
				t.Errorf("Expected role '%s', got '%s'", op.role, auditResult["role"])
			}

			if auditResult["success"] != op.success {
				t.Errorf("Expected success %v, got %v", op.success, auditResult["success"])
			}

			// Check for audit trail
			if auditTrail, ok := auditResult["audit_trail"].(string); !ok || auditTrail == "" {
				t.Errorf("Expected non-empty audit trail")
			}
		})
	}
}

// TestKeyRotationScheduling tests key rotation scheduling
func TestKeyRotationScheduling(t *testing.T) {
	cfg := &config.Config{
		Agent: config.AgentConfig{
			ID: "test-agent",
		},
		Security: config.SecurityConfig{
			KeyRotationDays: 7, // Short rotation for testing
			AuditLogging:    true,
			AuditPath:       "/tmp/test-rotation-audit.log",
		},
	}

	mockLogger := &MockAuditLogger{events: make([]*SecurityEvent, 0)}
	mockKeyStore := &MockKeyStore{keys: make(map[string]*SecureKeyData)}

	km, err := NewKeyManager(cfg, mockKeyStore, mockLogger)
	if err != nil {
		t.Fatalf("Failed to create key manager: %v", err)
	}
	defer func() {
		if km.rotationScheduler.running {
			km.rotationScheduler.Stop()
		}
	}()

	// Test that rotation jobs are scheduled for all roles
	roles := config.GetPredefinedRoles()
	expectedJobs := len(roles)

	if len(km.rotationScheduler.scheduledJobs) != expectedJobs {
		t.Errorf("Expected %d rotation jobs, got %d", expectedJobs, len(km.rotationScheduler.scheduledJobs))
	}

	// Test rotation policy is correctly set
	for _, job := range km.rotationScheduler.scheduledJobs {
		if job.Policy.RotationInterval != 7*24*time.Hour {
			t.Errorf("Expected rotation interval of 7 days, got %v", job.Policy.RotationInterval)
		}
		if !job.Policy.AutoRotate {
			t.Errorf("Expected auto-rotate to be enabled")
		}
	}
}

// Mock implementations for testing

type MockAuditLogger struct {
	events             []*SecurityEvent
	keyRotationEvents  []*KeyRotationEvent
}

func (m *MockAuditLogger) LogAccess(entry *AccessLogEntry) error {
	// Implementation for testing
	return nil
}

func (m *MockAuditLogger) LogKeyRotation(event *KeyRotationEvent) error {
	m.keyRotationEvents = append(m.keyRotationEvents, event)
	return nil
}

func (m *MockAuditLogger) LogSecurityEvent(event *SecurityEvent) error {
	m.events = append(m.events, event)
	return nil
}

func (m *MockAuditLogger) GetAuditTrail(criteria *AuditCriteria) ([]*AuditEvent, error) {
	return []*AuditEvent{}, nil
}

type MockKeyStore struct {
	keys     map[string]*SecureKeyData
	metadata []*KeyMetadata
}

func (m *MockKeyStore) StoreKey(keyID string, keyData *SecureKeyData) error {
	m.keys[keyID] = keyData
	return nil
}

func (m *MockKeyStore) RetrieveKey(keyID string) (*SecureKeyData, error) {
	if key, exists := m.keys[keyID]; exists {
		return key, nil
	}
	return nil, fmt.Errorf("key not found: %s", keyID)
}

func (m *MockKeyStore) DeleteKey(keyID string) error {
	delete(m.keys, keyID)
	return nil
}

func (m *MockKeyStore) ListKeys(filter *KeyFilter) ([]*KeyMetadata, error) {
	return m.metadata, nil
}

func (m *MockKeyStore) BackupKeys(criteria *BackupCriteria) (*KeyBackup, error) {
	return &KeyBackup{}, nil
}

func (m *MockKeyStore) RestoreKeys(backup *KeyBackup) error {
	return nil
}

// Test helper functions

func checkStoreAccessPolicyTest(role string) error {
	roles := config.GetPredefinedRoles()
	if _, exists := roles[role]; !exists {
		return fmt.Errorf("unknown creator role: %s", role)
	}
	
	roleData := roles[role]
	if roleData.AuthorityLevel == config.AuthorityReadOnly {
		return fmt.Errorf("role %s has read-only authority and cannot store content", role)
	}
	
	return nil
}

func checkRetrieveAccessPolicyTest(role string) error {
	roles := config.GetPredefinedRoles()
	if _, exists := roles[role]; !exists {
		return fmt.Errorf("unknown current role: %s", role)
	}
	
	return nil
}

func checkAnnounceAccessPolicyTest(role string) error {
	roles := config.GetPredefinedRoles()
	if _, exists := roles[role]; !exists {
		return fmt.Errorf("unknown current role: %s", role)
	}
	
	roleData := roles[role]
	if roleData.AuthorityLevel == config.AuthorityReadOnly || roleData.AuthorityLevel == config.AuthoritySuggestion {
		return fmt.Errorf("role %s lacks authority to announce content", role)
	}
	
	return nil
}

func simulateAuditOperation(cfg *config.Config, operation, ucxlAddress, role string, success bool, errorMsg string) map[string]interface{} {
	if !cfg.Security.AuditLogging || cfg.Security.AuditPath == "" {
		return nil
	}
	
	auditEntry := map[string]interface{}{
		"timestamp":     time.Now(),
		"operation":     operation,
		"node_id":       "test-node",
		"ucxl_address":  ucxlAddress,
		"role":          role,
		"success":       success,
		"error_message": errorMsg,
		"audit_trail":   fmt.Sprintf("DHT-%s-%s-%d", operation, ucxlAddress, time.Now().Unix()),
	}
	
	return auditEntry
}

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

// Benchmarks for security operations

func BenchmarkSecurityPolicyCheck(b *testing.B) {
	roles := []string{"admin", "backend_developer", "frontend_developer", "security_expert"}
	
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		role := roles[i%len(roles)]
		checkStoreAccessPolicyTest(role)
	}
}

func BenchmarkAuditLogging(b *testing.B) {
	cfg := &config.Config{
		Agent: config.AgentConfig{ID: "bench-agent", Role: "backend_developer"},
		Security: config.SecurityConfig{AuditLogging: true, AuditPath: "/tmp/bench-audit.log"},
	}
	
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		simulateAuditOperation(cfg, "store", "test:address:bench:task", "backend_developer", true, "")
	}
}

func BenchmarkKeyRotationCheck(b *testing.B) {
	cfg := &config.Config{
		Agent: config.AgentConfig{ID: "bench-agent"},
		Security: config.SecurityConfig{KeyRotationDays: 90, AuditLogging: true},
	}
	
	mockLogger := &MockAuditLogger{events: make([]*SecurityEvent, 0)}
	mockKeyStore := &MockKeyStore{
		keys:     make(map[string]*SecureKeyData),
		metadata: []*KeyMetadata{},
	}
	
	// Add some test keys
	for i := 0; i < 10; i++ {
		keyMeta := &KeyMetadata{
			KeyID:     fmt.Sprintf("bench-key-%d", i),
			KeyType:   "age-x25519",
			RoleID:    "backend_developer",
			CreatedAt: time.Now().Add(-time.Duration(i*10) * 24 * time.Hour),
			Status:    KeyStatusActive,
		}
		mockKeyStore.metadata = append(mockKeyStore.metadata, keyMeta)
	}
	
	km, err := NewKeyManager(cfg, mockKeyStore, mockLogger)
	if err != nil {
		b.Fatalf("Failed to create key manager: %v", err)
	}
	defer func() {
		if km.rotationScheduler.running {
			km.rotationScheduler.Stop()
		}
	}()
	
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		km.checkKeysForRotation()
	}
}