// Package crypto provides Age encryption implementation for role-based content security in CHORUS.
//
// This package implements the cryptographic foundation for CHORUS Phase 2B, enabling:
// - Role-based content encryption using Age (https://age-encryption.org)
// - Hierarchical access control based on agent authority levels
// - Multi-recipient encryption for shared content
// - Secure key management and validation
//
// The Age encryption system ensures that UCXL content is encrypted before storage
// in the distributed DHT, with access control enforced through role-based key distribution.
//
// Architecture Overview:
//   - Each role has an Age key pair (public/private)
//   - Content is encrypted for specific roles based on creator's authority
//   - Higher authority roles can decrypt lower authority content
//   - Admin roles can decrypt all content in the system
//
// Security Model:
//   - X25519 elliptic curve cryptography (Age standard)
//   - Per-role key pairs for access segmentation
//   - Authority hierarchy prevents privilege escalation
//   - Shamir secret sharing for admin key distribution (see shamir.go)
//
// Cross-references:
//   - pkg/config/roles.go: Role definitions and authority levels
//   - pkg/dht/encrypted_storage.go: Encrypted DHT storage implementation
//   - pkg/ucxl/decision_publisher.go: Decision publishing with encryption
//   - docs/ARCHITECTURE.md: Complete system architecture
//   - docs/SECURITY.md: Security model and threat analysis
package crypto

import (
	"bytes"
	"fmt"
	"io"
	"strings"

	"filippo.io/age"           // Modern, secure encryption library
	"chorus/pkg/config"
)

// AgeCrypto handles Age encryption for role-based content security.
//
// This is the primary interface for encrypting and decrypting UCXL content
// based on CHORUS role hierarchies. It provides methods to:
// - Encrypt content for specific roles or multiple roles
// - Decrypt content using the current agent's role key
// - Validate Age key formats and generate new key pairs
// - Determine decryption permissions based on role authority
//
// Usage Example:
//   crypto := NewAgeCrypto(config)
//   encrypted, err := crypto.EncryptForRole(content, "backend_developer")
//   decrypted, err := crypto.DecryptWithRole(encrypted)
//
// Thread Safety: AgeCrypto is safe for concurrent use across goroutines.
type AgeCrypto struct {
	config *config.Config // CHORUS configuration containing role definitions
}

// NewAgeCrypto creates a new Age crypto handler for role-based encryption.
//
// Parameters:
//   cfg: CHORUS configuration containing role definitions and agent settings
//
// Returns:
//   *AgeCrypto: Configured crypto handler ready for encryption/decryption
//
// The returned AgeCrypto instance will use the role definitions from the
// provided configuration to determine encryption permissions and key access.
//
// Cross-references:
//   - pkg/config/config.go: Configuration structure
//   - pkg/config/roles.go: Role definitions and authority levels
func NewAgeCrypto(cfg *config.Config) *AgeCrypto {
	return &AgeCrypto{
		config: cfg,
	}
}

// GenerateAgeKeyPair generates a new Age X25519 key pair for role-based encryption.
//
// This function creates cryptographically secure Age key pairs suitable for
// role-based content encryption. Each role in CHORUS should have its own key pair
// to enable proper access control and content segmentation.
//
// Returns:
//   *config.AgeKeyPair: Structure containing both public and private keys
//   error: Any error during key generation
//
// Key Format:
//   - Private key: "AGE-SECRET-KEY-1..." (Age standard format)
//   - Public key: "age1..." (Age recipient format)
//
// Security Notes:
//   - Uses X25519 elliptic curve cryptography
//   - Keys are cryptographically random using crypto/rand
//   - Private keys should be stored securely and never shared
//   - Public keys can be distributed freely for encryption
//
// Usage:
//   keyPair, err := GenerateAgeKeyPair()
//   if err != nil {
//       return fmt.Errorf("key generation failed: %w", err)
//   }
//   // Store keyPair.PrivateKey securely
//   // Distribute keyPair.PublicKey for encryption
//
// Cross-references:
//   - pkg/config/roles.go: AgeKeyPair structure definition
//   - docs/SECURITY.md: Key management best practices
//   - pkg/crypto/shamir.go: Admin key distribution via secret sharing
func GenerateAgeKeyPair() (*config.AgeKeyPair, error) {
	// Generate X25519 identity using Age's secure random generation
	identity, err := age.GenerateX25519Identity()
	if err != nil {
		return nil, fmt.Errorf("failed to generate Age identity: %w", err)
	}
	
	// Extract public and private key strings in Age format
	return &config.AgeKeyPair{
		PublicKey:  identity.Recipient().String(), // "age1..." format for recipients
		PrivateKey: identity.String(),             // "AGE-SECRET-KEY-1..." format
	}, nil
}

// ParseAgeIdentity parses an Age private key string into a usable identity.
//
// This function converts a private key string (AGE-SECRET-KEY-1...) into
// an Age identity that can be used for decryption operations.
//
// Parameters:
//   privateKey: Age private key string in standard format
//
// Returns:
//   age.Identity: Parsed identity for decryption operations
//   error: Parsing error if key format is invalid
//
// Key Format Requirements:
//   - Must start with "AGE-SECRET-KEY-1"
//   - Must be properly formatted X25519 private key
//   - Must be base64-encoded as per Age specification
//
// Cross-references:
//   - DecryptWithPrivateKey(): Uses parsed identities for decryption
//   - ValidateAgeKey(): Validates key format before parsing
func ParseAgeIdentity(privateKey string) (age.Identity, error) {
	return age.ParseX25519Identity(privateKey)
}

// ParseAgeRecipient parses an Age public key string into a recipient.
//
// This function converts a public key string (age1...) into an Age recipient
// that can be used for encryption operations.
//
// Parameters:
//   publicKey: Age public key string in recipient format
//
// Returns:
//   age.Recipient: Parsed recipient for encryption operations
//   error: Parsing error if key format is invalid
//
// Key Format Requirements:
//   - Must start with "age1"
//   - Must be properly formatted X25519 public key
//   - Must be base32-encoded as per Age specification
//
// Cross-references:
//   - EncryptForRole(): Uses parsed recipients for encryption
//   - ValidateAgeKey(): Validates key format before parsing
func ParseAgeRecipient(publicKey string) (age.Recipient, error) {
	return age.ParseX25519Recipient(publicKey)
}

// EncryptForRole encrypts content for a specific role using Age encryption
func (ac *AgeCrypto) EncryptForRole(content []byte, roleName string) ([]byte, error) {
	// Get role definition
	roles := config.GetPredefinedRoles()
	role, exists := roles[roleName]
	if !exists {
		return nil, fmt.Errorf("role '%s' not found", roleName)
	}
	
	// Check if role has Age keys configured
	if role.AgeKeys.PublicKey == "" {
		return nil, fmt.Errorf("role '%s' has no Age public key configured", roleName)
	}
	
	// Parse the recipient
	recipient, err := ParseAgeRecipient(role.AgeKeys.PublicKey)
	if err != nil {
		return nil, fmt.Errorf("failed to parse Age recipient for role '%s': %w", roleName, err)
	}
	
	// Encrypt the content
	out := &bytes.Buffer{}
	w, err := age.Encrypt(out, recipient)
	if err != nil {
		return nil, fmt.Errorf("failed to create Age encryptor: %w", err)
	}
	
	if _, err := w.Write(content); err != nil {
		return nil, fmt.Errorf("failed to write content to Age encryptor: %w", err)
	}
	
	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("failed to close Age encryptor: %w", err)
	}
	
	return out.Bytes(), nil
}

// EncryptForMultipleRoles encrypts content for multiple roles
func (ac *AgeCrypto) EncryptForMultipleRoles(content []byte, roleNames []string) ([]byte, error) {
	if len(roleNames) == 0 {
		return nil, fmt.Errorf("no roles specified")
	}
	
	var recipients []age.Recipient
	roles := config.GetPredefinedRoles()
	
	// Collect all recipients
	for _, roleName := range roleNames {
		role, exists := roles[roleName]
		if !exists {
			return nil, fmt.Errorf("role '%s' not found", roleName)
		}
		
		if role.AgeKeys.PublicKey == "" {
			return nil, fmt.Errorf("role '%s' has no Age public key configured", roleName)
		}
		
		recipient, err := ParseAgeRecipient(role.AgeKeys.PublicKey)
		if err != nil {
			return nil, fmt.Errorf("failed to parse Age recipient for role '%s': %w", roleName, err)
		}
		
		recipients = append(recipients, recipient)
	}
	
	// Encrypt for all recipients
	out := &bytes.Buffer{}
	w, err := age.Encrypt(out, recipients...)
	if err != nil {
		return nil, fmt.Errorf("failed to create Age encryptor: %w", err)
	}
	
	if _, err := w.Write(content); err != nil {
		return nil, fmt.Errorf("failed to write content to Age encryptor: %w", err)
	}
	
	if err := w.Close(); err != nil {
		return nil, fmt.Errorf("failed to close Age encryptor: %w", err)
	}
	
	return out.Bytes(), nil
}

// DecryptWithRole decrypts content using the current agent's role key
func (ac *AgeCrypto) DecryptWithRole(encryptedContent []byte) ([]byte, error) {
	if ac.config.Agent.Role == "" {
		return nil, fmt.Errorf("no role configured for current agent")
	}
	
	// Get current role's private key
	roles := config.GetPredefinedRoles()
	role, exists := roles[ac.config.Agent.Role]
	if !exists {
		return nil, fmt.Errorf("current role '%s' not found", ac.config.Agent.Role)
	}
	
	if role.AgeKeys.PrivateKey == "" {
		return nil, fmt.Errorf("current role '%s' has no Age private key configured", ac.config.Agent.Role)
	}
	
	return ac.DecryptWithPrivateKey(encryptedContent, role.AgeKeys.PrivateKey)
}

// DecryptWithPrivateKey decrypts content using a specific private key
func (ac *AgeCrypto) DecryptWithPrivateKey(encryptedContent []byte, privateKey string) ([]byte, error) {
	// Parse the identity
	identity, err := ParseAgeIdentity(privateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to parse Age identity: %w", err)
	}
	
	// Decrypt the content
	in := bytes.NewReader(encryptedContent)
	r, err := age.Decrypt(in, identity)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt content: %w", err)
	}
	
	out := &bytes.Buffer{}
	if _, err := io.Copy(out, r); err != nil {
		return nil, fmt.Errorf("failed to read decrypted content: %w", err)
	}
	
	return out.Bytes(), nil
}

// CanDecryptContent checks if current role can decrypt content encrypted for a target role
func (ac *AgeCrypto) CanDecryptContent(targetRole string) (bool, error) {
	return ac.config.CanDecryptRole(targetRole)
}

// GetDecryptableRoles returns list of roles current agent can decrypt
func (ac *AgeCrypto) GetDecryptableRoles() ([]string, error) {
	if ac.config.Agent.Role == "" {
		return nil, fmt.Errorf("no role configured")
	}
	
	roles := config.GetPredefinedRoles()
	currentRole, exists := roles[ac.config.Agent.Role]
	if !exists {
		return nil, fmt.Errorf("current role '%s' not found", ac.config.Agent.Role)
	}
	
	return currentRole.CanDecrypt, nil
}

// EncryptUCXLContent encrypts UCXL content based on creator's authority level
func (ac *AgeCrypto) EncryptUCXLContent(content []byte, creatorRole string) ([]byte, error) {
	// Get roles that should be able to decrypt this content
	decryptableRoles, err := ac.getDecryptableRolesForCreator(creatorRole)
	if err != nil {
		return nil, fmt.Errorf("failed to determine decryptable roles: %w", err)
	}
	
	// Encrypt for all decryptable roles
	return ac.EncryptForMultipleRoles(content, decryptableRoles)
}

// getDecryptableRolesForCreator determines which roles should be able to decrypt content from a creator
func (ac *AgeCrypto) getDecryptableRolesForCreator(creatorRole string) ([]string, error) {
	roles := config.GetPredefinedRoles()
	_, exists := roles[creatorRole]
	if !exists {
		return nil, fmt.Errorf("creator role '%s' not found", creatorRole)
	}
	
	// Start with the creator role itself
	decryptableRoles := []string{creatorRole}
	
	// Add all roles that have higher or equal authority and can decrypt this role
	for roleName, role := range roles {
		// Skip the creator role (already added)
		if roleName == creatorRole {
			continue
		}
		
		// Check if this role can decrypt the creator's content
		for _, decryptableRole := range role.CanDecrypt {
			if decryptableRole == creatorRole || decryptableRole == "*" {
				// Add this role to the list if not already present
				if !contains(decryptableRoles, roleName) {
					decryptableRoles = append(decryptableRoles, roleName)
				}
				break
			}
		}
	}
	
	return decryptableRoles, nil
}

// ValidateAgeKey validates an Age key format
func ValidateAgeKey(key string, isPrivate bool) error {
	if key == "" {
		return fmt.Errorf("key cannot be empty")
	}
	
	if isPrivate {
		// Validate private key format
		if !strings.HasPrefix(key, "AGE-SECRET-KEY-") {
			return fmt.Errorf("invalid Age private key format")
		}
		
		// Try to parse it
		_, err := ParseAgeIdentity(key)
		if err != nil {
			return fmt.Errorf("failed to parse Age private key: %w", err)
		}
	} else {
		// Validate public key format
		if !strings.HasPrefix(key, "age1") {
			return fmt.Errorf("invalid Age public key format")
		}
		
		// Try to parse it
		_, err := ParseAgeRecipient(key)
		if err != nil {
			return fmt.Errorf("failed to parse Age public key: %w", err)
		}
	}
	
	return nil
}

// GenerateRoleKeys generates Age key pairs for all roles that don't have them
func GenerateRoleKeys() (map[string]*config.AgeKeyPair, error) {
	roleKeys := make(map[string]*config.AgeKeyPair)
	roles := config.GetPredefinedRoles()
	
	for roleName, role := range roles {
		// Skip if role already has keys
		if role.AgeKeys.PublicKey != "" && role.AgeKeys.PrivateKey != "" {
			continue
		}
		
		// Generate new key pair
		keyPair, err := GenerateAgeKeyPair()
		if err != nil {
			return nil, fmt.Errorf("failed to generate keys for role '%s': %w", roleName, err)
		}
		
		roleKeys[roleName] = keyPair
	}
	
	return roleKeys, nil
}

// TestAgeEncryption tests Age encryption/decryption with sample data
func TestAgeEncryption() error {
	// Generate test key pair
	keyPair, err := GenerateAgeKeyPair()
	if err != nil {
		return fmt.Errorf("failed to generate test key pair: %w", err)
	}
	
	// Test content
	testContent := []byte("This is a test UCXL decision node content for Age encryption")
	
	// Parse recipient and identity
	recipient, err := ParseAgeRecipient(keyPair.PublicKey)
	if err != nil {
		return fmt.Errorf("failed to parse test recipient: %w", err)
	}
	
	identity, err := ParseAgeIdentity(keyPair.PrivateKey)
	if err != nil {
		return fmt.Errorf("failed to parse test identity: %w", err)
	}
	
	// Encrypt
	out := &bytes.Buffer{}
	w, err := age.Encrypt(out, recipient)
	if err != nil {
		return fmt.Errorf("failed to create test encryptor: %w", err)
	}
	
	if _, err := w.Write(testContent); err != nil {
		return fmt.Errorf("failed to write test content: %w", err)
	}
	
	if err := w.Close(); err != nil {
		return fmt.Errorf("failed to close test encryptor: %w", err)
	}
	
	encryptedContent := out.Bytes()
	
	// Decrypt
	in := bytes.NewReader(encryptedContent)
	r, err := age.Decrypt(in, identity)
	if err != nil {
		return fmt.Errorf("failed to decrypt test content: %w", err)
	}
	
	decryptedBuffer := &bytes.Buffer{}
	if _, err := io.Copy(decryptedBuffer, r); err != nil {
		return fmt.Errorf("failed to read decrypted test content: %w", err)
	}
	
	decryptedContent := decryptedBuffer.Bytes()
	
	// Verify
	if !bytes.Equal(testContent, decryptedContent) {
		return fmt.Errorf("test failed: decrypted content doesn't match original")
	}
	
	return nil
}

// contains checks if a string slice contains a value
func contains(slice []string, value string) bool {
	for _, item := range slice {
		if item == value {
			return true
		}
	}
	return false
}