Trixy3/trixy_core/ml/voice_recognition/model_loader.py

"""
Secure model loader for Voice Recognition system.

This module provides secure loading of PyTorch models with password protection
and metadata management. Supports .pth files as encrypted ZIP archives with
embedded metadata.
"""

import os
import json
import logging
import hashlib
import zipfile
import tempfile
import shutil
from typing import Dict, Any, Optional, List, Union, Tuple
from dataclasses import dataclass, asdict
from pathlib import Path
from datetime import datetime
import torch
import torch.nn as nn
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64

from .model_architectures import create_model_from_config, ModelArchitectureType

# Configure logging
logger = logging.getLogger(__name__)


class VRModelLoadError(Exception):
    """Exception raised when model loading fails."""
    pass


@dataclass
class VRModelMetadata:
    """Metadata for voice recognition models."""
    # Model information
    model_name: str
    model_version: str
    architecture_type: str
    embedding_dim: int
    num_speakers: int
    
    # Speaker mapping
    speaker_mapping: Dict[str, str]  # speaker_id -> speaker_name
    speaker_ids: List[str]  # ordered list of speaker IDs
    
    # Training information
    training_date: str
    training_duration_hours: float
    training_epochs: int
    training_samples: int
    final_accuracy: float
    final_loss: float
    
    # Audio configuration
    sample_rate: int
    audio_length_seconds: float
    feature_type: str
    feature_dim: int
    
    # Model configuration
    model_config: Dict[str, Any]
    
    # System information
    trained_on_device: str
    trained_by: str
    creation_timestamp: str
    file_hash: str
    
    # Compatibility information
    min_trixy_version: str = "1.0.0"
    model_format: str = "pytorch"
    
    def __post_init__(self):
        """Validate metadata after initialization."""
        if self.num_speakers != len(self.speaker_mapping):
            logger.warning(f"Mismatch: num_speakers={self.num_speakers}, speaker_mapping={len(self.speaker_mapping)}")
        
        if self.architecture_type not in [arch.value for arch in ModelArchitectureType]:
            raise ValueError(f"Unsupported architecture type: {self.architecture_type}")
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert metadata to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'VRModelMetadata':
        """Create metadata from dictionary."""
        return cls(**data)
    
    def validate_compatibility(self, required_speakers: Optional[List[str]] = None) -> List[str]:
        """
        Validate model compatibility.
        
        Args:
            required_speakers: List of required speaker names
            
        Returns:
            List of compatibility warnings
        """
        warnings = []
        
        # Check speaker compatibility
        if required_speakers:
            available_speakers = set(self.speaker_mapping.values())
            required_set = set(required_speakers)
            
            missing_speakers = required_set - available_speakers
            if missing_speakers:
                warnings.append(f"Missing speakers: {list(missing_speakers)}")
            
            extra_speakers = available_speakers - required_set
            if extra_speakers:
                warnings.append(f"Extra speakers in model: {list(extra_speakers)}")
        
        # Check feature compatibility
        if self.feature_dim not in [40, 80]:  # Common dimensions
            warnings.append(f"Unusual feature dimension: {self.feature_dim}")
        
        if self.sample_rate != 16000:
            warnings.append(f"Non-standard sample rate: {self.sample_rate}Hz")
        
        return warnings


class VRModelLoader:
    """Basic model loader without encryption."""
    
    def __init__(self, device: str = "auto"):
        """
        Initialize model loader.
        
        Args:
            device: Device to load model on
        """
        if device == "auto":
            if torch.cuda.is_available():
                self.device = "cuda"
            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                self.device = "mps"
            else:
                self.device = "cpu"
        else:
            self.device = device
        
        logger.info(f"Model loader initialized with device: {self.device}")
    
    def load_metadata(self, password: str = "") -> VRModelMetadata:
        """
        Load only metadata from model file without loading the full model.

        Args:
            password: Password for encrypted models (currently unused for basic loader)

        Returns:
            Model metadata
        """
        # This is called during initialization with model_path passed to __init__
        # We need to extract model_path from self or accept it as parameter
        # For now, this is a placeholder that raises an error
        raise NotImplementedError("load_metadata requires model_path - use SecureVRModelLoader instead")

    def load_model(self, model_path: str, password: str = "") -> Tuple[nn.Module, VRModelMetadata]:
        """
        Load model and metadata from file (ZIP or direct .pth).

        Args:
            model_path: Path to model file
            password: Password for encrypted ZIP files

        Returns:
            Tuple of (model, metadata)
        """
        if not os.path.exists(model_path):
            raise VRModelLoadError(f"Model file not found: {model_path}")

        logger.info(f"Loading model from: {model_path}")

        # Check if it's a ZIP file (trainer format)
        try:
            with zipfile.ZipFile(model_path, 'r') as zip_file:
                # It's a ZIP file - load from metafile.json + model.pth
                return self._load_from_zip(zip_file, password)
        except zipfile.BadZipFile:
            # Not a ZIP - try loading as direct PyTorch checkpoint
            return self._load_from_checkpoint(model_path)

    def _load_from_zip(self, zip_file: zipfile.ZipFile, password: str) -> Tuple[nn.Module, VRModelMetadata]:
        """Load model from ZIP archive with metafile.json."""
        import io

        # Set password if provided
        if password:
            zip_file.setpassword(password.encode('utf-8'))

        # Load metadata
        try:
            with zip_file.open('metafile.json') as meta_file:
                trainer_metadata = json.load(meta_file)
                metadata = self._convert_trainer_metadata(trainer_metadata)
        except KeyError:
            raise VRModelLoadError("metafile.json not found in model ZIP")

        # Load model state
        try:
            with zip_file.open('model.pth') as model_file:
                model_bytes = model_file.read()
                state_dict = torch.load(
                    io.BytesIO(model_bytes),
                    map_location=self.device,
                    weights_only=True
                )
        except KeyError:
            raise VRModelLoadError("model.pth not found in model ZIP")

        # Create model
        model = self._create_model_from_metadata(metadata)
        model.load_state_dict(state_dict, strict=False)
        model.to(self.device)
        model.eval()

        logger.info(f"Model loaded from ZIP: {metadata.model_name}")
        return model, metadata

    def _load_from_checkpoint(self, model_path: str) -> Tuple[nn.Module, VRModelMetadata]:
        """Load model from direct PyTorch checkpoint file."""
        try:
            checkpoint = torch.load(model_path, map_location=self.device)
        except Exception as e:
            raise VRModelLoadError(f"Failed to load PyTorch model: {e}")

        # Extract metadata
        if 'metadata' not in checkpoint:
            raise VRModelLoadError("Model file missing metadata")

        metadata = VRModelMetadata.from_dict(checkpoint['metadata'])

        # Create model architecture
        model = self._create_model_from_metadata(metadata)

        # Load model weights
        if 'model_state_dict' in checkpoint:
            try:
                model.load_state_dict(checkpoint['model_state_dict'])
            except Exception as e:
                raise VRModelLoadError(f"Failed to load model weights: {e}")
        else:
            raise VRModelLoadError("Model file missing state_dict")

        model.to(self.device)
        model.eval()

        logger.info(f"Model loaded from checkpoint: {metadata.model_name}")
        return model, metadata

    def _convert_trainer_metadata(self, trainer_meta: Dict[str, Any]) -> VRModelMetadata:
        """Convert trainer metafile.json format to VRModelMetadata format."""
        # Extract nested data
        training_stats = trainer_meta.get('training_statistics', {})
        audio_cfg = trainer_meta.get('audio_config', {})
        arch_info = trainer_meta.get('architecture_info', {})
        hardware = trainer_meta.get('hardware_info', {})
        model_specific = trainer_meta.get('model_specific', {})
        dataset_info = trainer_meta.get('dataset_info', {})

        # Build speaker mapping from class_names or model_specific
        speaker_ids = dataset_info.get('class_names', [])
        if not speaker_ids and 'speakers' in model_specific:
            speaker_ids = model_specific['speakers']

        speaker_mapping = {str(i): name for i, name in enumerate(speaker_ids)}
        num_speakers = len(speaker_ids)

        # Map to VRModelMetadata fields
        # Convert architecture name to lowercase to match enum values
        architecture_name = arch_info.get('architecture_name', 'ecapa_tdnn')
        if isinstance(architecture_name, str):
            architecture_name = architecture_name.lower()

        return VRModelMetadata(
            model_name=trainer_meta.get('model_name', 'unknown'),
            model_version=trainer_meta.get('version', '1.0.0'),
            architecture_type=architecture_name,
            embedding_dim=arch_info.get('embedding_dim', 192),
            num_speakers=num_speakers,

            speaker_mapping=speaker_mapping,
            speaker_ids=speaker_ids,

            training_date=trainer_meta.get('creation_timestamp', ''),
            training_duration_hours=training_stats.get('training_duration_seconds', 0.0) / 3600.0,
            training_epochs=training_stats.get('total_epochs', 0),
            training_samples=dataset_info.get('total_samples', 0),
            final_accuracy=training_stats.get('best_val_accuracy', 0.0),
            final_loss=training_stats.get('best_val_loss', 0.0),

            sample_rate=audio_cfg.get('sample_rate', 16000),
            audio_length_seconds=audio_cfg.get('audio_length_seconds', 3.0),
            feature_type=audio_cfg.get('feature_type', 'mel_spectrogram'),
            feature_dim=audio_cfg.get('n_mels', 80),

            model_config=arch_info,

            trained_on_device=hardware.get('gpu_names', ['cpu'])[0] if hardware.get('gpu_names') else 'cpu',
            trained_by=trainer_meta.get('author', 'Unknown'),
            creation_timestamp=trainer_meta.get('creation_timestamp', ''),
            file_hash='',

            min_trixy_version='1.0.0',
            model_format='pytorch'
        )
    
    def _create_model_from_metadata(self, metadata: VRModelMetadata) -> nn.Module:
        """Create model architecture from metadata."""
        # Extract model_config and remove keys that we're passing explicitly
        model_config = metadata.model_config.copy() if metadata.model_config else {}
        # Remove keys that might conflict with explicit parameters
        model_config.pop('architecture_type', None)
        model_config.pop('input_dim', None)
        model_config.pop('embedding_dim', None)
        model_config.pop('num_speakers', None)

        return create_model_from_config(
            architecture_type=metadata.architecture_type,
            input_dim=metadata.feature_dim,
            embedding_dim=metadata.embedding_dim,
            num_speakers=metadata.num_speakers,
            **model_config
        )
    
    def save_model(self, 
                   model: nn.Module, 
                   metadata: VRModelMetadata, 
                   save_path: str,
                   optimizer: Optional[torch.optim.Optimizer] = None) -> None:
        """
        Save model with metadata.
        
        Args:
            model: Model to save
            metadata: Model metadata
            save_path: Path to save model
            optimizer: Optional optimizer state
        """
        # Prepare checkpoint
        checkpoint = {
            'model_state_dict': model.state_dict(),
            'metadata': metadata.to_dict(),
            'save_timestamp': datetime.now().isoformat()
        }
        
        if optimizer is not None:
            checkpoint['optimizer_state_dict'] = optimizer.state_dict()
        
        # Create directory if needed
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        
        # Save model
        try:
            torch.save(checkpoint, save_path)
            logger.info(f"Model saved to: {save_path}")
        except Exception as e:
            raise VRModelLoadError(f"Failed to save model: {e}")


class SecureVRModelLoader(VRModelLoader):
    """Secure model loader with password protection and ZIP support."""
    
    def __init__(self, device: str = "auto"):
        super().__init__(device)
        self.temp_dir = None
    
    def load_model(self, 
                   model_path: str, 
                   password: Optional[str] = None) -> Tuple[nn.Module, VRModelMetadata]:
        """
        Load encrypted model from ZIP archive.
        
        Args:
            model_path: Path to encrypted model file
            password: Password for decryption
            
        Returns:
            Tuple of (model, metadata)
        """
        if not os.path.exists(model_path):
            raise VRModelLoadError(f"Model file not found: {model_path}")
        
        logger.info(f"Loading secure model from: {model_path}")
        
        # Create temporary directory
        self.temp_dir = tempfile.mkdtemp(prefix="trixy_vr_model_")
        
        try:
            # Extract and decrypt model
            if model_path.endswith('.zip') or self._is_zip_file(model_path):
                model_data, metadata_data = self._load_from_zip(model_path, password)
            elif model_path.endswith('.pth'):
                # Try to load as encrypted ZIP (PTH files can be ZIP archives)
                try:
                    model_data, metadata_data = self._load_from_zip(model_path, password)
                except:
                    # Fall back to regular PyTorch loading
                    return super().load_model(model_path)
            else:
                raise VRModelLoadError(f"Unsupported model format: {model_path}")
            
            # Parse metadata
            metadata = VRModelMetadata.from_dict(metadata_data)
            
            # Create model architecture
            model = self._create_model_from_metadata(metadata)
            
            # Load model weights
            try:
                model.load_state_dict(model_data['model_state_dict'])
            except Exception as e:
                raise VRModelLoadError(f"Failed to load model weights: {e}")
            
            model.to(self.device)
            model.eval()
            
            logger.info(f"Secure model loaded successfully: {metadata.model_name}")
            return model, metadata
            
        finally:
            # Clean up temporary directory
            if self.temp_dir and os.path.exists(self.temp_dir):
                shutil.rmtree(self.temp_dir)
                self.temp_dir = None
    
    def _is_zip_file(self, file_path: str) -> bool:
        """Check if file is a ZIP archive."""
        try:
            with zipfile.ZipFile(file_path, 'r') as zf:
                return True
        except:
            return False
    
    def _load_from_zip(self, zip_path: str, password: Optional[str] = None) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """Load model and metadata from ZIP archive."""
        try:
            with zipfile.ZipFile(zip_path, 'r') as zf:
                # List files in archive
                file_list = zf.namelist()
                logger.debug(f"Files in archive: {file_list}")
                
                # Find metadata file
                metadata_file = None
                for filename in file_list:
                    if filename.endswith('metadata.json') or filename == 'metadata.json':
                        metadata_file = filename
                        break
                
                if not metadata_file:
                    raise VRModelLoadError("No metadata.json found in model archive")
                
                # Find model file
                model_file = None
                for filename in file_list:
                    if filename.endswith('.pth') or filename.endswith('.pt'):
                        model_file = filename
                        break
                
                if not model_file:
                    raise VRModelLoadError("No model file (.pth/.pt) found in archive")
                
                # Extract files
                if password:
                    # Decrypt files
                    metadata_data = self._decrypt_file_content(zf.read(metadata_file), password)
                    model_content = self._decrypt_file_content(zf.read(model_file), password)
                else:
                    metadata_data = zf.read(metadata_file).decode('utf-8')
                    model_content = zf.read(model_file)
                
                # Parse metadata
                metadata_dict = json.loads(metadata_data)
                
                # Load model from bytes
                model_temp_path = os.path.join(self.temp_dir, 'model.pth')
                with open(model_temp_path, 'wb') as f:
                    f.write(model_content)
                
                model_data = torch.load(model_temp_path, map_location=self.device)
                
                return model_data, metadata_dict
                
        except zipfile.BadZipFile:
            raise VRModelLoadError("Invalid ZIP archive")
        except Exception as e:
            raise VRModelLoadError(f"Failed to load from ZIP: {e}")
    
    def _decrypt_file_content(self, encrypted_content: bytes, password: str) -> str:
        """Decrypt file content using password."""
        try:
            # Generate key from password
            key = self._generate_key_from_password(password)
            
            # Decrypt content
            fernet = Fernet(key)
            decrypted_content = fernet.decrypt(encrypted_content)
            
            return decrypted_content.decode('utf-8')
            
        except Exception as e:
            raise VRModelLoadError(f"Failed to decrypt content: {e}")
    
    def _generate_key_from_password(self, password: str) -> bytes:
        """Generate encryption key from password."""
        # Use a fixed salt for consistency (in production, should be stored)
        salt = b'trixy_voice_recognition_salt_2024'
        
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        return key
    
    def save_secure_model(self,
                         model: nn.Module,
                         metadata: VRModelMetadata,
                         save_path: str,
                         password: Optional[str] = None,
                         optimizer: Optional[torch.optim.Optimizer] = None) -> None:
        """
        Save model as encrypted ZIP archive.
        
        Args:
            model: Model to save
            metadata: Model metadata
            save_path: Path to save encrypted model
            password: Password for encryption
            optimizer: Optional optimizer state
        """
        # Create temporary directory
        temp_dir = tempfile.mkdtemp(prefix="trixy_vr_save_")
        
        try:
            # Prepare model checkpoint
            checkpoint = {
                'model_state_dict': model.state_dict(),
                'save_timestamp': datetime.now().isoformat()
            }
            
            if optimizer is not None:
                checkpoint['optimizer_state_dict'] = optimizer.state_dict()
            
            # Save model to temporary file
            model_temp_path = os.path.join(temp_dir, 'model.pth')
            torch.save(checkpoint, model_temp_path)
            
            # Prepare metadata
            metadata_dict = metadata.to_dict()
            metadata_json = json.dumps(metadata_dict, indent=2)
            
            # Create ZIP archive
            Path(save_path).parent.mkdir(parents=True, exist_ok=True)
            
            with zipfile.ZipFile(save_path, 'w', zipfile.ZIP_DEFLATED) as zf:
                if password:
                    # Encrypt files
                    encrypted_metadata = self._encrypt_content(metadata_json, password)
                    
                    with open(model_temp_path, 'rb') as f:
                        model_content = f.read()
                    encrypted_model = self._encrypt_content(model_content, password)
                    
                    # Add encrypted files to archive
                    zf.writestr('metadata.json', encrypted_metadata)
                    zf.writestr('model.pth', encrypted_model)
                else:
                    # Add unencrypted files
                    zf.writestr('metadata.json', metadata_json)
                    zf.write(model_temp_path, 'model.pth')
                
                # Add readme
                readme_content = f"""
Trixy Voice Recognition Model
============================

Model Name: {metadata.model_name}
Architecture: {metadata.architecture_type}
Speakers: {metadata.num_speakers}
Created: {metadata.creation_timestamp}

This is a secure voice recognition model for the Trixy voice assistant.
Load using the SecureVRModelLoader class.
"""
                zf.writestr('README.txt', readme_content)
            
            logger.info(f"Secure model saved to: {save_path}")
            
        finally:
            # Clean up
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
    
    def _encrypt_content(self, content: Union[str, bytes], password: str) -> bytes:
        """Encrypt content using password."""
        if isinstance(content, str):
            content = content.encode('utf-8')
        
        # Generate key from password
        key = self._generate_key_from_password(password)
        
        # Encrypt content
        fernet = Fernet(key)
        encrypted_content = fernet.encrypt(content)
        
        return encrypted_content


def validate_vr_model_file(model_path: str, password: Optional[str] = None) -> Dict[str, Any]:
    """
    Validate voice recognition model file.
    
    Args:
        model_path: Path to model file
        password: Password if model is encrypted
        
    Returns:
        Validation results
    """
    results = {
        'is_valid': False,
        'file_exists': False,
        'file_size_mb': 0.0,
        'is_encrypted': False,
        'has_metadata': False,
        'metadata': None,
        'errors': [],
        'warnings': []
    }
    
    try:
        # Check file existence
        if not os.path.exists(model_path):
            results['errors'].append(f"File not found: {model_path}")
            return results
        
        results['file_exists'] = True
        results['file_size_mb'] = os.path.getsize(model_path) / (1024 * 1024)
        
        # Try to load model
        loader = SecureVRModelLoader()
        
        try:
            model, metadata = loader.load_model(model_path, password)
            results['has_metadata'] = True
            results['metadata'] = metadata.to_dict()
            results['is_valid'] = True
            
            # Validate metadata
            validation_warnings = metadata.validate_compatibility()
            results['warnings'].extend(validation_warnings)
            
        except VRModelLoadError as e:
            if "decrypt" in str(e).lower():
                results['is_encrypted'] = True
                results['errors'].append("Model is encrypted - password required or incorrect")
            else:
                results['errors'].append(f"Model loading failed: {e}")
        
    except Exception as e:
        results['errors'].append(f"Validation error: {e}")
    
    return results


def create_vr_model_metadata(model_name: str,
                            architecture_type: str,
                            embedding_dim: int,
                            speaker_mapping: Dict[str, str],
                            **kwargs) -> VRModelMetadata:
    """
    Create model metadata with defaults.
    
    Args:
        model_name: Name of the model
        architecture_type: Type of architecture
        embedding_dim: Embedding dimension
        speaker_mapping: Mapping of speaker IDs to names
        **kwargs: Additional metadata fields
        
    Returns:
        Model metadata
    """
    speaker_ids = list(speaker_mapping.keys())
    
    metadata = VRModelMetadata(
        model_name=model_name,
        model_version=kwargs.get('model_version', '1.0.0'),
        architecture_type=architecture_type,
        embedding_dim=embedding_dim,
        num_speakers=len(speaker_mapping),
        speaker_mapping=speaker_mapping,
        speaker_ids=speaker_ids,
        training_date=kwargs.get('training_date', datetime.now().strftime('%Y-%m-%d')),
        training_duration_hours=kwargs.get('training_duration_hours', 0.0),
        training_epochs=kwargs.get('training_epochs', 0),
        training_samples=kwargs.get('training_samples', 0),
        final_accuracy=kwargs.get('final_accuracy', 0.0),
        final_loss=kwargs.get('final_loss', 0.0),
        sample_rate=kwargs.get('sample_rate', 16000),
        audio_length_seconds=kwargs.get('audio_length_seconds', 3.0),
        feature_type=kwargs.get('feature_type', 'log_mel'),
        feature_dim=kwargs.get('feature_dim', 40),
        model_config=kwargs.get('model_config', {}),
        trained_on_device=kwargs.get('trained_on_device', 'unknown'),
        trained_by=kwargs.get('trained_by', 'trixy_trainer'),
        creation_timestamp=datetime.now().isoformat(),
        file_hash=kwargs.get('file_hash', ''),
        min_trixy_version=kwargs.get('min_trixy_version', '1.0.0'),
        model_format=kwargs.get('model_format', 'pytorch')
    )
    
    return metadata


# Example usage and testing
def test_model_loader():
    """Test model loading functionality."""
    from .model_architectures import ECADATDNNModel
    
    # Create sample model
    model = ECADATDNNModel(
        input_dim=40,
        embedding_dim=192,
        num_speakers=4
    )
    
    # Create sample metadata
    speaker_mapping = {
        "spk_001": "Alice",
        "spk_002": "Bob", 
        "spk_003": "Charlie",
        "spk_004": "Diana"
    }
    
    metadata = create_vr_model_metadata(
        model_name="test_ecapa_model",
        architecture_type="ecapa_tdnn",
        embedding_dim=192,
        speaker_mapping=speaker_mapping,
        training_epochs=100,
        final_accuracy=0.95
    )
    
    # Test basic saving and loading
    temp_path = "/tmp/test_vr_model.pth"
    
    loader = VRModelLoader()
    loader.save_model(model, metadata, temp_path)
    
    loaded_model, loaded_metadata = loader.load_model(temp_path)
    
    print(f"Model loaded: {loaded_metadata.model_name}")
    print(f"Speakers: {loaded_metadata.speaker_mapping}")
    
    # Test secure saving and loading
    secure_loader = SecureVRModelLoader()
    secure_path = "/tmp/test_vr_model_secure.pth"
    password = "test_password_123"
    
    secure_loader.save_secure_model(model, metadata, secure_path, password)
    
    loaded_secure_model, loaded_secure_metadata = secure_loader.load_model(secure_path, password)
    
    print(f"Secure model loaded: {loaded_secure_metadata.model_name}")
    
    # Clean up
    if os.path.exists(temp_path):
        os.remove(temp_path)
    if os.path.exists(secure_path):
        os.remove(secure_path)


if __name__ == "__main__":
    test_model_loader()