Trixy3/trainer/model_formats.py

"""
Model format handlers for different file formats (.pth, .pt, .onnx).

This module provides handlers for saving and loading models in various formats,
including password protection and metadata integration as specified in the
project requirements.
"""

import os
import json
import zipfile
import tempfile
import hashlib
import io
from abc import ABC, abstractmethod
from typing import Dict, Any, Optional, Union, BinaryIO
from pathlib import Path
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
import logging

try:
    import onnx
    import onnxruntime as ort
    ONNX_AVAILABLE = True
except ImportError:
    ONNX_AVAILABLE = False
    onnx = None
    ort = None


class ModelFormatError(Exception):
    """Exception raised for model format related errors."""
    pass


class EncryptionManager:
    """Handles encryption and decryption for password-protected models."""
    
    @staticmethod
    def generate_key_from_password(password: str, salt: bytes = None) -> bytes:
        """Generate encryption key from password using PBKDF2."""
        if salt is None:
            salt = os.urandom(16)
        
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        return key, salt
    
    @staticmethod
    def encrypt_data(data: bytes, password: str) -> bytes:
        """Encrypt data with password."""
        key, salt = EncryptionManager.generate_key_from_password(password)
        f = Fernet(key)
        encrypted_data = f.encrypt(data)
        # Prepend salt to encrypted data
        return salt + encrypted_data
    
    @staticmethod
    def decrypt_data(encrypted_data: bytes, password: str) -> bytes:
        """Decrypt data with password."""
        # Extract salt from the beginning
        salt = encrypted_data[:16]
        encrypted_content = encrypted_data[16:]
        
        key, _ = EncryptionManager.generate_key_from_password(password, salt)
        f = Fernet(key)
        return f.decrypt(encrypted_content)


class ModelFormatHandler(ABC):
    """Abstract base class for model format handlers."""
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        self.logger = logger or logging.getLogger(__name__)
    
    @abstractmethod
    def save_model(self, model: nn.Module, filepath: str, metadata: Dict[str, Any], 
                   password: Optional[str] = None, **kwargs) -> bool:
        """Save model with metadata and optional password protection."""
        pass
    
    @abstractmethod
    def load_model(self, filepath: str, password: Optional[str] = None, 
                   map_location: Optional[str] = None) -> tuple:
        """Load model and return (model_state_dict, metadata)."""
        pass
    
    @abstractmethod
    def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]:
        """Extract metadata from model file without loading the full model."""
        pass
    
    @abstractmethod
    def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], 
                       password: Optional[str] = None) -> bool:
        """Update metadata in existing model file."""
        pass
    
    def _generate_file_hash(self, filepath: str) -> str:
        """Generate SHA256 hash of the file."""
        hash_sha256 = hashlib.sha256()
        with open(filepath, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_sha256.update(chunk)
        return hash_sha256.hexdigest()


class PTHHandler(ModelFormatHandler):
    """Handler for .pth files (ZIP archives with JSON metadata)."""
    
    def save_model(self, model: nn.Module, filepath: str, metadata: Dict[str, Any], 
                   password: Optional[str] = None, **kwargs) -> bool:
        """
        Save model as .pth file (ZIP archive) with metadata.json.
        
        Args:
            model: PyTorch model to save
            filepath: Path to save the model
            metadata: Metadata dictionary to include
            password: Optional password for encryption
            **kwargs: Additional arguments
            
        Returns:
            bool: True if successful
        """
        try:
            # Ensure directory exists
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            
            # Create temporary directory for building the archive
            with tempfile.TemporaryDirectory() as temp_dir:
                # Save model state dict
                model_path = os.path.join(temp_dir, "model.pth")
                torch.save(model.state_dict(), model_path)
                
                # Save metadata
                metadata_path = os.path.join(temp_dir, "metafile.json")
                metadata['file_format'] = 'pth'
                metadata['pytorch_version'] = torch.__version__
                metadata['file_hash'] = None  # Will be calculated after creation
                
                with open(metadata_path, 'w', encoding='utf-8') as f:
                    json.dump(metadata, f, indent=2, ensure_ascii=False)
                
                # Create ZIP archive
                with zipfile.ZipFile(filepath, 'w', zipfile.ZIP_DEFLATED) as zipf:
                    zipf.write(model_path, "model.pth")
                    zipf.write(metadata_path, "metafile.json")
                    
                    # Add any additional files from kwargs
                    for key, value in kwargs.items():
                        if key.startswith('add_file_'):
                            file_path = value.get('path')
                            archive_name = value.get('name', os.path.basename(file_path))
                            if os.path.exists(file_path):
                                zipf.write(file_path, archive_name)
            
            # Apply password protection if requested
            if password:
                self._encrypt_file(filepath, password)
            
            # Update file hash in metadata
            file_hash = self._generate_file_hash(filepath)
            self._update_file_hash(filepath, file_hash, password)
            
            self.logger.info(f"Model saved successfully: {filepath}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to save model: {str(e)}")
            return False
    
    def load_model(self, filepath: str, password: Optional[str] = None, 
                   map_location: Optional[str] = None) -> tuple:
        """
        Load model from .pth file.
        
        Returns:
            tuple: (model_state_dict, metadata)
        """
        try:
            # Handle password protection
            if password:
                filepath = self._decrypt_file(filepath, password)
            
            with zipfile.ZipFile(filepath, 'r') as zipf:
                # Load metadata
                with zipf.open("metafile.json") as f:
                    metadata = json.load(f)
                
                # Load model
                with zipf.open("model.pth") as f:
                    model_data = f.read()
                    model_state_dict = torch.load(
                        io.BytesIO(model_data), 
                        map_location=map_location
                    )
            
            self.logger.info(f"Model loaded successfully: {filepath}")
            return model_state_dict, metadata
            
        except Exception as e:
            self.logger.error(f"Failed to load model: {str(e)}")
            raise ModelFormatError(f"Failed to load model: {str(e)}")
    
    def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]:
        """Extract metadata from .pth file without loading the model."""
        try:
            # Handle password protection
            if password:
                filepath = self._decrypt_file(filepath, password)
            
            with zipfile.ZipFile(filepath, 'r') as zipf:
                with zipf.open("metafile.json") as f:
                    metadata = json.load(f)
            
            return metadata
            
        except Exception as e:
            self.logger.error(f"Failed to extract metadata: {str(e)}")
            return {}
    
    def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], 
                       password: Optional[str] = None) -> bool:
        """Update metadata in existing .pth file."""
        try:
            # Load existing data
            model_state_dict, existing_metadata = self.load_model(filepath, password)
            
            # Merge metadata
            existing_metadata.update(new_metadata)
            
            # Create temporary model for saving
            temp_model = nn.Module()
            temp_model.load_state_dict = lambda: None
            temp_model.state_dict = lambda: model_state_dict
            
            # Save with updated metadata
            return self.save_model(temp_model, filepath, existing_metadata, password)
            
        except Exception as e:
            self.logger.error(f"Failed to update metadata: {str(e)}")
            return False
    
    def _encrypt_file(self, filepath: str, password: str):
        """Encrypt file in place."""
        with open(filepath, 'rb') as f:
            data = f.read()
        
        encrypted_data = EncryptionManager.encrypt_data(data, password)
        
        with open(filepath, 'wb') as f:
            f.write(encrypted_data)
    
    def _decrypt_file(self, filepath: str, password: str) -> str:
        """Decrypt file to temporary location and return temp path."""
        with open(filepath, 'rb') as f:
            encrypted_data = f.read()
        
        try:
            decrypted_data = EncryptionManager.decrypt_data(encrypted_data, password)
        except Exception:
            raise ModelFormatError("Invalid password or corrupted file")
        
        # Save to temporary file
        temp_fd, temp_path = tempfile.mkstemp(suffix='.pth')
        with os.fdopen(temp_fd, 'wb') as f:
            f.write(decrypted_data)
        
        return temp_path
    
    def _update_file_hash(self, filepath: str, file_hash: str, password: Optional[str] = None):
        """Update file hash in metadata."""
        try:
            # Load and update metadata
            metadata = self.extract_metadata(filepath, password)
            metadata['file_hash'] = file_hash
            self.update_metadata(filepath, metadata, password)
        except Exception as e:
            self.logger.warning(f"Failed to update file hash: {str(e)}")


class PTHandler(ModelFormatHandler):
    """Handler for .pt files using torch.package."""
    
    def save_model(self, model: nn.Module, filepath: str, metadata: Dict[str, Any], 
                   password: Optional[str] = None, **kwargs) -> bool:
        """
        Save model as .pt file using torch.package.
        
        Args:
            model: PyTorch model to save
            filepath: Path to save the model
            metadata: Metadata dictionary to include
            password: Optional password for encryption
            **kwargs: Additional arguments
            
        Returns:
            bool: True if successful
        """
        try:
            # Ensure directory exists
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            
            # Create package
            with torch.package.PackageExporter(filepath) as exp:
                # Add model
                exp.intern('**')
                exp.save_pickle('model', 'model.pkl', model.state_dict())
                
                # Add metadata
                metadata['file_format'] = 'pt'
                metadata['pytorch_version'] = torch.__version__
                metadata['package_version'] = '1.0'
                
                # Save metadata as JSON string in package
                exp.save_text('metadata', 'metafile.json', 
                             json.dumps(metadata, indent=2, ensure_ascii=False))
                
                # Add custom resources if specified
                for key, value in kwargs.items():
                    if key.startswith('add_resource_'):
                        resource_name = key.replace('add_resource_', '')
                        if isinstance(value, str):
                            exp.save_text('resources', f'{resource_name}.txt', value)
                        elif isinstance(value, (dict, list)):
                            exp.save_text('resources', f'{resource_name}.json', 
                                        json.dumps(value, indent=2))
            
            # Apply password protection if requested
            if password:
                self._encrypt_file(filepath, password)
            
            # Update file hash
            file_hash = self._generate_file_hash(filepath)
            self._update_file_hash(filepath, file_hash, password)
            
            self.logger.info(f"Model saved successfully: {filepath}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to save model: {str(e)}")
            return False
    
    def load_model(self, filepath: str, password: Optional[str] = None, 
                   map_location: Optional[str] = None) -> tuple:
        """
        Load model from .pt file.
        
        Returns:
            tuple: (model_state_dict, metadata)
        """
        try:
            # Handle password protection
            if password:
                filepath = self._decrypt_file(filepath, password)
            
            # Load package
            imp = torch.package.PackageImporter(filepath)
            
            # Load model state dict
            model_state_dict = imp.load_pickle('model', 'model.pkl', 
                                             map_location=map_location)
            
            # Load metadata
            metadata_json = imp.load_text('metadata', 'metafile.json')
            metadata = json.loads(metadata_json)
            
            self.logger.info(f"Model loaded successfully: {filepath}")
            return model_state_dict, metadata
            
        except Exception as e:
            self.logger.error(f"Failed to load model: {str(e)}")
            raise ModelFormatError(f"Failed to load model: {str(e)}")
    
    def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]:
        """Extract metadata from .pt file without loading the model."""
        try:
            # Handle password protection
            if password:
                filepath = self._decrypt_file(filepath, password)
            
            imp = torch.package.PackageImporter(filepath)
            metadata_json = imp.load_text('metadata', 'metafile.json')
            metadata = json.loads(metadata_json)
            
            return metadata
            
        except Exception as e:
            self.logger.error(f"Failed to extract metadata: {str(e)}")
            return {}
    
    def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], 
                       password: Optional[str] = None) -> bool:
        """Update metadata in existing .pt file."""
        try:
            # Load existing data
            model_state_dict, existing_metadata = self.load_model(filepath, password)
            
            # Merge metadata
            existing_metadata.update(new_metadata)
            
            # Create temporary model for saving
            temp_model = nn.Module()
            temp_model.load_state_dict = lambda: None
            temp_model.state_dict = lambda: model_state_dict
            
            # Save with updated metadata
            return self.save_model(temp_model, filepath, existing_metadata, password)
            
        except Exception as e:
            self.logger.error(f"Failed to update metadata: {str(e)}")
            return False
    
    def _encrypt_file(self, filepath: str, password: str):
        """Encrypt file in place."""
        with open(filepath, 'rb') as f:
            data = f.read()
        
        encrypted_data = EncryptionManager.encrypt_data(data, password)
        
        with open(filepath, 'wb') as f:
            f.write(encrypted_data)
    
    def _decrypt_file(self, filepath: str, password: str) -> str:
        """Decrypt file to temporary location and return temp path."""
        with open(filepath, 'rb') as f:
            encrypted_data = f.read()
        
        try:
            decrypted_data = EncryptionManager.decrypt_data(encrypted_data, password)
        except Exception:
            raise ModelFormatError("Invalid password or corrupted file")
        
        # Save to temporary file
        temp_fd, temp_path = tempfile.mkstemp(suffix='.pt')
        with os.fdopen(temp_fd, 'wb') as f:
            f.write(decrypted_data)
        
        return temp_path
    
    def _update_file_hash(self, filepath: str, file_hash: str, password: Optional[str] = None):
        """Update file hash in metadata."""
        try:
            metadata = self.extract_metadata(filepath, password)
            metadata['file_hash'] = file_hash
            self.update_metadata(filepath, metadata, password)
        except Exception as e:
            self.logger.warning(f"Failed to update file hash: {str(e)}")


class ONNXHandler(ModelFormatHandler):
    """Handler for .onnx files with metadata_props."""
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        super().__init__(logger)
        if not ONNX_AVAILABLE:
            raise ImportError("ONNX not available. Install with: pip install onnx onnxruntime")
    
    def save_model(self, model: Union[nn.Module, str], filepath: str, metadata: Dict[str, Any], 
                   password: Optional[str] = None, **kwargs) -> bool:
        """
        Save model as .onnx file with metadata_props.
        
        Args:
            model: PyTorch model or path to existing ONNX model
            filepath: Path to save the model
            metadata: Metadata dictionary to include
            password: Optional password for encryption
            **kwargs: Additional arguments for ONNX export
            
        Returns:
            bool: True if successful
        """
        try:
            # Ensure directory exists
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            
            if isinstance(model, str):
                # Load existing ONNX model
                onnx_model = onnx.load(model)
            else:
                # Export PyTorch model to ONNX
                dummy_input = kwargs.get('dummy_input')
                if dummy_input is None:
                    raise ValueError("dummy_input required for PyTorch to ONNX conversion")
                
                # Create temporary ONNX file
                temp_fd, temp_onnx_path = tempfile.mkstemp(suffix='.onnx')
                os.close(temp_fd)
                
                torch.onnx.export(
                    model,
                    dummy_input,
                    temp_onnx_path,
                    export_params=True,
                    opset_version=kwargs.get('opset_version', 11),
                    do_constant_folding=True,
                    input_names=kwargs.get('input_names', ['input']),
                    output_names=kwargs.get('output_names', ['output']),
                    dynamic_axes=kwargs.get('dynamic_axes', {}),
                    verbose=False
                )
                
                onnx_model = onnx.load(temp_onnx_path)
                os.unlink(temp_onnx_path)
            
            # Add metadata to model
            metadata['file_format'] = 'onnx'
            metadata['onnx_version'] = onnx.__version__
            
            # Clear existing metadata_props
            del onnx_model.metadata_props[:]
            
            # Add new metadata as key-value pairs
            for key, value in metadata.items():
                prop = onnx_model.metadata_props.add()
                prop.key = key
                if isinstance(value, (dict, list)):
                    prop.value = json.dumps(value, ensure_ascii=False)
                else:
                    prop.value = str(value)
            
            # Save model
            temp_filepath = filepath
            if password:
                temp_fd, temp_filepath = tempfile.mkstemp(suffix='.onnx')
                os.close(temp_fd)
            
            onnx.save(onnx_model, temp_filepath)
            
            # Apply password protection if requested
            if password:
                self._encrypt_file(temp_filepath, password)
                # Move encrypted file to final location
                with open(temp_filepath, 'rb') as src, open(filepath, 'wb') as dst:
                    dst.write(src.read())
                os.unlink(temp_filepath)
            
            # Update file hash
            file_hash = self._generate_file_hash(filepath)
            self._update_file_hash(filepath, file_hash, password)
            
            self.logger.info(f"Model saved successfully: {filepath}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to save model: {str(e)}")
            return False
    
    def load_model(self, filepath: str, password: Optional[str] = None, 
                   map_location: Optional[str] = None) -> tuple:
        """
        Load model from .onnx file.
        
        Returns:
            tuple: (onnx_model, metadata)
        """
        try:
            # Handle password protection
            if password:
                filepath = self._decrypt_file(filepath, password)
            
            # Load ONNX model
            onnx_model = onnx.load(filepath)
            
            # Extract metadata from metadata_props
            metadata = {}
            for prop in onnx_model.metadata_props:
                try:
                    # Try to parse as JSON first
                    metadata[prop.key] = json.loads(prop.value)
                except json.JSONDecodeError:
                    # If not JSON, store as string
                    metadata[prop.key] = prop.value
            
            self.logger.info(f"Model loaded successfully: {filepath}")
            return onnx_model, metadata
            
        except Exception as e:
            self.logger.error(f"Failed to load model: {str(e)}")
            raise ModelFormatError(f"Failed to load model: {str(e)}")
    
    def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]:
        """Extract metadata from .onnx file without loading the full model."""
        try:
            # Handle password protection
            if password:
                filepath = self._decrypt_file(filepath, password)
            
            # Load model to extract metadata
            onnx_model = onnx.load(filepath)
            
            # Extract metadata
            metadata = {}
            for prop in onnx_model.metadata_props:
                try:
                    metadata[prop.key] = json.loads(prop.value)
                except json.JSONDecodeError:
                    metadata[prop.key] = prop.value
            
            return metadata
            
        except Exception as e:
            self.logger.error(f"Failed to extract metadata: {str(e)}")
            return {}
    
    def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], 
                       password: Optional[str] = None) -> bool:
        """Update metadata in existing .onnx file."""
        try:
            # Load existing model and metadata
            onnx_model, existing_metadata = self.load_model(filepath, password)
            
            # Merge metadata
            existing_metadata.update(new_metadata)
            
            # Save with updated metadata
            return self.save_model(onnx_model, filepath, existing_metadata, password)
            
        except Exception as e:
            self.logger.error(f"Failed to update metadata: {str(e)}")
            return False
    
    def _encrypt_file(self, filepath: str, password: str):
        """Encrypt file in place."""
        with open(filepath, 'rb') as f:
            data = f.read()
        
        encrypted_data = EncryptionManager.encrypt_data(data, password)
        
        with open(filepath, 'wb') as f:
            f.write(encrypted_data)
    
    def _decrypt_file(self, filepath: str, password: str) -> str:
        """Decrypt file to temporary location and return temp path."""
        with open(filepath, 'rb') as f:
            encrypted_data = f.read()
        
        try:
            decrypted_data = EncryptionManager.decrypt_data(encrypted_data, password)
        except Exception:
            raise ModelFormatError("Invalid password or corrupted file")
        
        # Save to temporary file
        temp_fd, temp_path = tempfile.mkstemp(suffix='.onnx')
        with os.fdopen(temp_fd, 'wb') as f:
            f.write(decrypted_data)
        
        return temp_path
    
    def _update_file_hash(self, filepath: str, file_hash: str, password: Optional[str] = None):
        """Update file hash in metadata."""
        try:
            metadata = self.extract_metadata(filepath, password)
            metadata['file_hash'] = file_hash
            self.update_metadata(filepath, metadata, password)
        except Exception as e:
            self.logger.warning(f"Failed to update file hash: {str(e)}")


class ModelFormatManager:
    """Manager class for handling different model formats."""
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        self.logger = logger or logging.getLogger(__name__)
        self.handlers = {
            '.pth': PTHHandler(logger),
            '.pt': PTHandler(logger),
        }
        
        # Add ONNX handler if available
        if ONNX_AVAILABLE:
            self.handlers['.onnx'] = ONNXHandler(logger)
        else:
            self.logger.warning("ONNX support not available")
    
    def get_handler(self, filepath: str) -> ModelFormatHandler:
        """Get appropriate handler for file format."""
        ext = Path(filepath).suffix.lower()
        if ext not in self.handlers:
            raise ModelFormatError(f"Unsupported format: {ext}")
        return self.handlers[ext]
    
    def save_model(self, model: Union[nn.Module, str], filepath: str, 
                   metadata: Dict[str, Any], password: Optional[str] = None, 
                   **kwargs) -> bool:
        """Save model using appropriate handler."""
        handler = self.get_handler(filepath)
        return handler.save_model(model, filepath, metadata, password, **kwargs)
    
    def load_model(self, filepath: str, password: Optional[str] = None, 
                   map_location: Optional[str] = None) -> tuple:
        """Load model using appropriate handler."""
        handler = self.get_handler(filepath)
        return handler.load_model(filepath, password, map_location)
    
    def extract_metadata(self, filepath: str, password: Optional[str] = None) -> Dict[str, Any]:
        """Extract metadata using appropriate handler."""
        handler = self.get_handler(filepath)
        return handler.extract_metadata(filepath, password)
    
    def update_metadata(self, filepath: str, new_metadata: Dict[str, Any], 
                       password: Optional[str] = None) -> bool:
        """Update metadata using appropriate handler."""
        handler = self.get_handler(filepath)
        return handler.update_metadata(filepath, new_metadata, password)
    
    def get_supported_formats(self) -> List[str]:
        """Get list of supported file formats."""
        return list(self.handlers.keys())