"""
AI Identity - Ed25519 Keypair Management

Implements cryptographic identity for AI agents using Ed25519 signatures.
Provides keypair generation, secure storage, and identity verification.

Design Principles:
- Ed25519 for fast, secure signatures
- File system permissions (0600) for private key security
- JSON format for interoperability
- Portable identity (can be exported/imported)

Storage:
- Location: .empirica/identity/<ai_id>.key
- Format: JSON with hex-encoded keys
- Permissions: 0600 (read/write owner only)
"""

import os
import json
import logging
from pathlib import Path
from typing import Optional, Dict, Any
from datetime import datetime, UTC

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey
)
from cryptography.hazmat.primitives import serialization

logger = logging.getLogger(__name__)


class AIIdentity:
    """
    AI Identity with Ed25519 keypair
    
    Usage:
        # Create new identity
        identity = AIIdentity(ai_id="claude-code")
        identity.generate_keypair()
        identity.save_keypair()
        
        # Load existing identity
        identity = AIIdentity(ai_id="claude-code")
        identity.load_keypair()
        
        # Use for signing
        signature = identity.sign(message)
        verified = identity.verify(signature, message, public_key)
    """
    
    def __init__(self, ai_id: str, identity_dir: Optional[str] = None):
        """
        Initialize AI Identity
        
        Args:
            ai_id: AI identifier (e.g., "claude-code", "mini-agent")
            identity_dir: Custom identity storage directory
        """
        self.ai_id = ai_id
        self.identity_dir = Path(identity_dir or ".empirica/identity")
        self.private_key: Optional[Ed25519PrivateKey] = None
        self.public_key: Optional[Ed25519PublicKey] = None
        self.created_at: Optional[str] = None
        self.metadata: Dict[str, Any] = {}
        
    @property
    def keypair_path(self) -> Path:
        """Path to keypair file"""
        return self.identity_dir / f"{self.ai_id}.key"
    
    @property
    def public_key_path(self) -> Path:
        """Path to public key file (for distribution)"""
        return self.identity_dir / f"{self.ai_id}.pub"
    
    def generate_keypair(self) -> None:
        """
        Generate new Ed25519 keypair
        
        Raises:
            RuntimeError: If keypair already exists
        """
        if self.private_key is not None:
            raise RuntimeError(f"Keypair already exists for {self.ai_id}")
        
        # Generate Ed25519 keypair
        self.private_key = Ed25519PrivateKey.generate()
        self.public_key = self.private_key.public_key()
        self.created_at = datetime.now(UTC).isoformat()
        
        logger.info(f"✓ Generated Ed25519 keypair for {self.ai_id}")
    
    def save_keypair(self, overwrite: bool = False) -> None:
        """
        Save keypair to disk
        
        Args:
            overwrite: Allow overwriting existing keypair
            
        Raises:
            RuntimeError: If no keypair or file exists without overwrite
        """
        if self.private_key is None:
            raise RuntimeError("No keypair to save. Call generate_keypair() first.")
        
        if self.keypair_path.exists() and not overwrite:
            raise RuntimeError(
                f"Keypair already exists at {self.keypair_path}. "
                "Use overwrite=True to replace."
            )
        
        # Ensure directory exists
        self.identity_dir.mkdir(parents=True, exist_ok=True)
        
        # Serialize keys to bytes
        private_bytes = self.private_key.private_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PrivateFormat.Raw,
            encryption_algorithm=serialization.NoEncryption()
        )
        
        public_bytes = self.public_key.public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw
        )
        
        # Create keypair JSON
        keypair_data = {
            'ai_id': self.ai_id,
            'created_at': self.created_at,
            'private_key': private_bytes.hex(),
            'public_key': public_bytes.hex(),
            'metadata': self.metadata
        }
        
        # Write keypair file
        with open(self.keypair_path, 'w') as f:
            json.dump(keypair_data, f, indent=2)
        
        # Set restrictive permissions (0600)
        os.chmod(self.keypair_path, 0o600)
        
        # Write public key file (for distribution)
        public_data = {
            'ai_id': self.ai_id,
            'created_at': self.created_at,
            'public_key': public_bytes.hex(),
            'metadata': self.metadata
        }
        
        with open(self.public_key_path, 'w') as f:
            json.dump(public_data, f, indent=2)
        
        logger.info(f"✓ Saved keypair to {self.keypair_path}")
        logger.info(f"✓ Saved public key to {self.public_key_path}")
    
    def load_keypair(self) -> None:
        """
        Load keypair from disk
        
        Raises:
            FileNotFoundError: If keypair file doesn't exist
        """
        if not self.keypair_path.exists():
            raise FileNotFoundError(
                f"Keypair not found at {self.keypair_path}. "
                f"Create one with: empirica identity-create --ai-id {self.ai_id}"
            )
        
        # Load keypair JSON
        with open(self.keypair_path, 'r') as f:
            keypair_data = json.load(f)
        
        # Verify ai_id matches
        if keypair_data['ai_id'] != self.ai_id:
            raise ValueError(
                f"AI ID mismatch: file has {keypair_data['ai_id']}, "
                f"expected {self.ai_id}"
            )
        
        # Deserialize keys
        private_bytes = bytes.fromhex(keypair_data['private_key'])
        public_bytes = bytes.fromhex(keypair_data['public_key'])
        
        self.private_key = Ed25519PrivateKey.from_private_bytes(private_bytes)
        self.public_key = Ed25519PublicKey.from_public_bytes(public_bytes)
        self.created_at = keypair_data['created_at']
        self.metadata = keypair_data.get('metadata', {})
        
        logger.info(f"✓ Loaded keypair for {self.ai_id}")
    
    def sign(self, message: bytes) -> bytes:
        """
        Sign message with private key
        
        Args:
            message: Message to sign
            
        Returns:
            bytes: Signature
            
        Raises:
            RuntimeError: If no private key loaded
        """
        if self.private_key is None:
            raise RuntimeError("No private key loaded. Call load_keypair() first.")
        
        return self.private_key.sign(message)
    
    @staticmethod
    def verify(signature: bytes, message: bytes, public_key_bytes: bytes) -> bool:
        """
        Verify signature with public key
        
        Args:
            signature: Signature to verify
            message: Original message
            public_key_bytes: Public key bytes
            
        Returns:
            bool: True if signature valid
        """
        try:
            public_key = Ed25519PublicKey.from_public_bytes(public_key_bytes)
            public_key.verify(signature, message)
            return True
        except Exception:
            return False
    
    def public_key_hex(self) -> str:
        """Get public key as hex string"""
        if self.public_key is None:
            raise RuntimeError("No public key loaded")
        
        return self.public_key.public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw
        ).hex()
    
    def export_public_key(self) -> Dict[str, Any]:
        """
        Export public key for sharing
        
        Returns:
            Dict with ai_id, public_key, created_at
        """
        if self.public_key is None:
            raise RuntimeError("No public key loaded")
        
        return {
            'ai_id': self.ai_id,
            'created_at': self.created_at,
            'public_key': self.public_key_hex(),
            'metadata': self.metadata
        }


class IdentityManager:
    """
    Manage multiple AI identities
    
    Usage:
        manager = IdentityManager()
        
        # List identities
        identities = manager.list_identities()
        
        # Load identity
        identity = manager.load_identity("claude-code")
        
        # Create identity
        identity = manager.create_identity("new-agent")
    """
    
    def __init__(self, identity_dir: Optional[str] = None):
        """Initialize identity manager"""
        self.identity_dir = Path(identity_dir or ".empirica/identity")
    
    def list_identities(self) -> list[Dict[str, Any]]:
        """
        List all identities
        
        Returns:
            List of identity summaries
        """
        if not self.identity_dir.exists():
            return []
        
        identities = []
        
        for key_file in self.identity_dir.glob("*.key"):
            try:
                with open(key_file, 'r') as f:
                    data = json.load(f)
                
                identities.append({
                    'ai_id': data['ai_id'],
                    'created_at': data['created_at'],
                    'key_file': str(key_file),
                    'has_public_key': (self.identity_dir / f"{data['ai_id']}.pub").exists()
                })
            except Exception as e:
                logger.warning(f"Failed to load {key_file}: {e}")
        
        return identities
    
    def load_identity(self, ai_id: str) -> AIIdentity:
        """
        Load existing identity
        
        Args:
            ai_id: AI identifier
            
        Returns:
            AIIdentity: Loaded identity
        """
        identity = AIIdentity(ai_id, str(self.identity_dir))
        identity.load_keypair()
        return identity
    
    def create_identity(self, ai_id: str, overwrite: bool = False) -> AIIdentity:
        """
        Create new identity
        
        Args:
            ai_id: AI identifier
            overwrite: Allow overwriting existing identity
            
        Returns:
            AIIdentity: New identity
        """
        identity = AIIdentity(ai_id, str(self.identity_dir))
        identity.generate_keypair()
        identity.save_keypair(overwrite=overwrite)
        return identity
    
    def identity_exists(self, ai_id: str) -> bool:
        """Check if identity exists"""
        return (self.identity_dir / f"{ai_id}.key").exists()
