integrations.py

#!/usr/bin/env python3
"""
AI Model and Integration Management
Specialized handlers for AI/ML models, Docker integration, and development tools
"""

import os
import json
import shutil
import tempfile
import subprocess
from pathlib import Path
from typing import Dict, List, Optional, Any
import yaml
import requests
from datetime import datetime

class AIModelManager:
    """Manager for AI/ML model artifacts and metadata."""
    
    def __init__(self, storage_path: str):
        self.storage_path = Path(storage_path)
        self.models_path = self.storage_path / 'models'
        self.models_path.mkdir(exist_ok=True)
    
    def detect_model_type(self, file_path: str) -> Dict[str, Any]:
        """Detect AI model type and extract metadata."""
        file_path = Path(file_path)
        metadata = {
            'framework': 'unknown',
            'model_type': 'unknown',
            'input_shape': None,
            'output_shape': None,
            'parameters': None,
            'size_mb': file_path.stat().st_size / (1024 * 1024)
        }
        
        ext = file_path.suffix.lower()
        
        # PyTorch models
        if ext in ['.pt', '.pth']:
            metadata['framework'] = 'pytorch'
            try:
                import torch
                model = torch.load(file_path, map_location='cpu')
                if hasattr(model, 'state_dict'):
                    metadata['parameters'] = sum(p.numel() for p in model.parameters())
                metadata['model_type'] = 'pytorch_model'
            except (ImportError, Exception):
                pass
        
        # TensorFlow models
        elif ext in ['.pb', '.h5']:
            metadata['framework'] = 'tensorflow'
            if ext == '.h5':
                try:
                    import tensorflow as tf
                    model = tf.keras.models.load_model(file_path)
                    metadata['parameters'] = model.count_params()
                    metadata['model_type'] = 'keras_model'
                except (ImportError, Exception):
                    pass
            else:
                metadata['model_type'] = 'tensorflow_savedmodel'
        
        # ONNX models
        elif ext == '.onnx':
            metadata['framework'] = 'onnx'
            metadata['model_type'] = 'onnx_model'
            try:
                import onnx
                model = onnx.load(file_path)
                metadata['input_shape'] = [
                    [dim.dim_value for dim in inp.type.tensor_type.shape.dim]
                    for inp in model.graph.input
                ]
                metadata['output_shape'] = [
                    [dim.dim_value for dim in out.type.tensor_type.shape.dim]
                    for out in model.graph.output
                ]
            except (ImportError, Exception):
                pass
        
        # Hugging Face models
        elif file_path.name in ['pytorch_model.bin', 'model.safetensors']:
            metadata['framework'] = 'huggingface'
            metadata['model_type'] = 'transformers_model'
            
            # Look for config.json in same directory
            config_path = file_path.parent / 'config.json'
            if config_path.exists():
                try:
                    with open(config_path) as f:
                        config = json.load(f)
                    metadata['model_config'] = config
                    metadata['architecture'] = config.get('architectures', ['unknown'])[0]
                except Exception:
                    pass
        
        # Scikit-learn models
        elif ext == '.pkl':
            metadata['framework'] = 'scikit-learn'
            metadata['model_type'] = 'sklearn_model'
            try:
                import pickle
                with open(file_path, 'rb') as f:
                    model = pickle.load(f)
                metadata['model_class'] = type(model).__name__
            except (ImportError, Exception):
                pass
        
        return metadata
    
    def create_model_card(self, model_metadata: Dict[str, Any], 
                         name: str, version: str, description: str = "") -> str:
        """Create a model card for the AI model."""
        card = {
            'name': name,
            'version': version,
            'description': description,
            'created_at': datetime.utcnow().isoformat(),
            'framework': model_metadata.get('framework', 'unknown'),
            'model_type': model_metadata.get('model_type', 'unknown'),
            'size_mb': model_metadata.get('size_mb', 0),
            'parameters': model_metadata.get('parameters'),
            'input_shape': model_metadata.get('input_shape'),
            'output_shape': model_metadata.get('output_shape'),
            'usage': {
                'inference_code': self._generate_inference_code(model_metadata),
                'requirements': self._get_requirements(model_metadata)
            }
        }
        
        return yaml.dump(card, default_flow_style=False)
    
    def _generate_inference_code(self, metadata: Dict[str, Any]) -> str:
        """Generate sample inference code based on model type."""
        framework = metadata.get('framework', 'unknown')
        
        if framework == 'pytorch':
            return """
import torch

# Load model
model = torch.load('model.pt')
model.eval()

# Inference
with torch.no_grad():
    output = model(input_tensor)
"""
        elif framework == 'tensorflow':
            return """
import tensorflow as tf

# Load model
model = tf.keras.models.load_model('model.h5')

# Inference
predictions = model.predict(input_data)
"""
        elif framework == 'onnx':
            return """
import onnxruntime as ort

# Load model
session = ort.InferenceSession('model.onnx')

# Inference
outputs = session.run(None, {'input': input_data})
"""
        elif framework == 'huggingface':
            return """
from transformers import AutoModel, AutoTokenizer

# Load model and tokenizer
model = AutoModel.from_pretrained('./model')
tokenizer = AutoTokenizer.from_pretrained('./model')

# Inference
inputs = tokenizer(text, return_tensors='pt')
outputs = model(**inputs)
"""
        else:
            return "# Add your inference code here"
    
    def _get_requirements(self, metadata: Dict[str, Any]) -> List[str]:
        """Get requirements based on model framework."""
        framework = metadata.get('framework', 'unknown')
        
        requirements_map = {
            'pytorch': ['torch', 'torchvision'],
            'tensorflow': ['tensorflow'],
            'onnx': ['onnxruntime'],
            'huggingface': ['transformers', 'torch'],
            'scikit-learn': ['scikit-learn', 'numpy'],
        }
        
        return requirements_map.get(framework, [])

class DockerRegistryHandler:
    """Handler for Docker Registry HTTP API V2 compatibility."""
    
    def __init__(self, storage_path: str):
        self.storage_path = Path(storage_path)
        self.docker_path = self.storage_path / 'docker'
        self.docker_path.mkdir(exist_ok=True)
    
    def handle_manifest(self, name: str, reference: str) -> Dict[str, Any]:
        """Handle Docker manifest requests."""
        manifest_path = self.docker_path / name / 'manifests' / reference
        
        if manifest_path.exists():
            with open(manifest_path) as f:
                return json.load(f)
        else:
            raise FileNotFoundError(f"Manifest not found: {name}:{reference}")
    
    def handle_blob(self, name: str, digest: str) -> str:
        """Handle Docker blob requests."""
        blob_path = self.docker_path / name / 'blobs' / digest
        
        if blob_path.exists():
            return str(blob_path)
        else:
            raise FileNotFoundError(f"Blob not found: {digest}")
    
    def store_image_layers(self, image_path: str, name: str, tag: str) -> Dict[str, Any]:
        """Extract and store Docker image layers."""
        import tarfile
        
        repo_path = self.docker_path / name
        repo_path.mkdir(exist_ok=True)
        (repo_path / 'manifests').mkdir(exist_ok=True)
        (repo_path / 'blobs').mkdir(exist_ok=True)
        
        # Extract tar file
        with tempfile.TemporaryDirectory() as temp_dir:
            with tarfile.open(image_path, 'r') as tar:
                tar.extractall(temp_dir)
            
            # Process manifest.json
            manifest_file = Path(temp_dir) / 'manifest.json'
            if manifest_file.exists():
                with open(manifest_file) as f:
                    manifest = json.load(f)[0]
                
                # Store manifest
                manifest_path = repo_path / 'manifests' / tag
                with open(manifest_path, 'w') as f:
                    json.dump(manifest, f, indent=2)
                
                # Store layers as blobs
                for layer in manifest.get('Layers', []):
                    layer_path = Path(temp_dir) / layer
                    if layer_path.exists():
                        # Calculate digest
                        import hashlib
                        with open(layer_path, 'rb') as f:
                            digest = hashlib.sha256(f.read()).hexdigest()
                        
                        blob_path = repo_path / 'blobs' / f"sha256:{digest}"
                        shutil.copy2(layer_path, blob_path)
                
                return manifest
        
        raise ValueError("Invalid Docker image format")

class DevToolsIntegration:
    """Integration with development tools and CI/CD systems."""
    
    def __init__(self, registry_url: str, token: str):
        self.registry_url = registry_url
        self.token = token
        self.session = requests.Session()
        self.session.headers.update({'Authorization': f'Bearer {token}'})
    
    def generate_maven_pom(self, artifact: Dict[str, Any]) -> str:
        """Generate Maven POM snippet for Java artifacts."""
        return f"""
<dependency>
    <groupId>{artifact.get('group_id', 'com.company')}</groupId>
    <artifactId>{artifact['name']}</artifactId>
    <version>{artifact['version']}</version>
</dependency>
"""
    
    def generate_npm_config(self, registry_url: str) -> str:
        """Generate NPM configuration."""
        return f"""
# .npmrc
registry={registry_url}/npm/
//{registry_url.split('://')[-1]}/npm/:_authToken={self.token}
"""
    
    def generate_pip_config(self, registry_url: str) -> str:
        """Generate pip configuration."""
        return f"""
# pip.conf
[global]
index-url = {registry_url}/pypi/simple/
trusted-host = {registry_url.split('://')[-1]}
"""
    
    def generate_docker_config(self, registry_url: str) -> str:
        """Generate Docker daemon configuration."""
        return f"""
# /etc/docker/daemon.json
{{
    "insecure-registries": ["{registry_url.split('://')[-1]}"],
    "registry-mirrors": ["{registry_url}/v2/"]
}}
"""
    
    def generate_ci_pipeline(self, artifact_type: str, org: str, repo: str) -> str:
        """Generate CI/CD pipeline configuration."""
        if artifact_type == 'docker':
            return self._generate_docker_pipeline(org, repo)
        elif artifact_type == 'python':
            return self._generate_python_pipeline(org, repo)
        elif artifact_type == 'npm':
            return self._generate_npm_pipeline(org, repo)
        else:
            return self._generate_generic_pipeline(org, repo)
    
    def _generate_docker_pipeline(self, org: str, repo: str) -> str:
        """Generate Docker CI pipeline."""
        return f"""
# .github/workflows/docker.yml
name: Docker Build and Push

on:
  push:
    branches: [ main ]
  pull_request:
    branches: [ main ]

jobs:
  build:
    runs-on: ubuntu-latest
    
    steps:
    - uses: actions/checkout@v3
    
    - name: Build Docker image
      run: docker build -t ${{{{ github.repository }}}}:${{{{ github.sha }}}} .
    
    - name: Push to registry
      run: |
        echo "${{{{ secrets.REGISTRY_TOKEN }}}}" | docker login {self.registry_url} -u ${{{{ secrets.REGISTRY_USER }}}} --password-stdin
        docker tag ${{{{ github.repository }}}}:${{{{ github.sha }}}} {self.registry_url}/{org}/{repo}:${{{{ github.sha }}}}
        docker push {self.registry_url}/{org}/{repo}:${{{{ github.sha }}}}
"""
    
    def _generate_python_pipeline(self, org: str, repo: str) -> str:
        """Generate Python CI pipeline."""
        return f"""
# .github/workflows/python.yml
name: Python Package

on:
  push:
    branches: [ main ]
  pull_request:
    branches: [ main ]

jobs:
  build:
    runs-on: ubuntu-latest
    
    steps:
    - uses: actions/checkout@v3
    
    - name: Set up Python
      uses: actions/setup-python@v4
      with:
        python-version: '3.9'
    
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        pip install build twine
    
    - name: Build package
      run: python -m build
    
    - name: Upload to registry
      run: |
        python -m pip install artifact-registry-cli
        artifact-registry-cli upload dist/*.whl {org} {repo} --token ${{{{ secrets.REGISTRY_TOKEN }}}}
"""
    
    def _generate_npm_pipeline(self, org: str, repo: str) -> str:
        """Generate NPM CI pipeline."""
        return f"""
# .github/workflows/npm.yml
name: NPM Package

on:
  push:
    branches: [ main ]
  pull_request:
    branches: [ main ]

jobs:
  build:
    runs-on: ubuntu-latest
    
    steps:
    - uses: actions/checkout@v3
    
    - name: Set up Node.js
      uses: actions/setup-node@v3
      with:
        node-version: '16'
    
    - name: Install dependencies
      run: npm install
    
    - name: Build
      run: npm run build
    
    - name: Package
      run: npm pack
    
    - name: Upload to registry
      run: |
        pip install artifact-registry-cli
        artifact-registry-cli upload *.tgz {org} {repo} --token ${{{{ secrets.REGISTRY_TOKEN }}}}
"""
    
    def _generate_generic_pipeline(self, org: str, repo: str) -> str:
        """Generate generic CI pipeline."""
        return f"""
# .github/workflows/artifact.yml
name: Build and Upload Artifact

on:
  push:
    branches: [ main ]
  pull_request:
    branches: [ main ]

jobs:
  build:
    runs-on: ubuntu-latest
    
    steps:
    - uses: actions/checkout@v3
    
    - name: Build artifact
      run: |
        # Add your build commands here
        echo "Building artifact..."
    
    - name: Upload to registry
      run: |
        pip install artifact-registry-cli
        artifact-registry-cli upload ./artifact.bin {org} {repo} --token ${{{{ secrets.REGISTRY_TOKEN }}}}
"""

class SecurityScanner:
    """Security scanning for uploaded artifacts."""
    
    def __init__(self):
        self.scanners = []
    
    def scan_artifact(self, file_path: str, artifact_type: str) -> Dict[str, Any]:
        """Scan artifact for security issues."""
        results = {
            'scanned_at': datetime.utcnow().isoformat(),
            'vulnerabilities': [],
            'warnings': [],
            'info': []
        }
        
        # Basic file analysis
        file_size = os.path.getsize(file_path)
        if file_size > 1024 * 1024 * 1024:  # 1GB
            results['warnings'].append({
                'type': 'large_file',
                'message': f'Large file detected: {file_size / (1024*1024*1024):.2f} GB'
            })
        
        # Artifact-specific scans
        if artifact_type == 'docker':
            results.update(self._scan_docker_image(file_path))
        elif artifact_type in ['python', 'npm']:
            results.update(self._scan_package_dependencies(file_path, artifact_type))
        
        return results
    
    def _scan_docker_image(self, image_path: str) -> Dict[str, Any]:
        """Scan Docker image for vulnerabilities."""
        results = {'docker_scan': True}
        
        try:
            # Use trivy if available
            result = subprocess.run([
                'trivy', 'image', '--format', 'json', image_path
            ], capture_output=True, text=True, timeout=300)
            
            if result.returncode == 0:
                scan_data = json.loads(result.stdout)
                results['trivy_scan'] = scan_data
        except (subprocess.TimeoutExpired, FileNotFoundError, json.JSONDecodeError):
            results['warnings'] = results.get('warnings', [])
            results['warnings'].append({
                'type': 'scan_failed',
                'message': 'Docker security scan failed or trivy not available'
            })
        
        return results
    
    def _scan_package_dependencies(self, package_path: str, package_type: str) -> Dict[str, Any]:
        """Scan package dependencies for known vulnerabilities."""
        results = {'dependency_scan': True}
        
        # This would integrate with vulnerability databases
        # For now, just placeholder
        results['info'].append({
            'type': 'dependency_check',
            'message': f'Dependency scan completed for {package_type} package'
        })
        
        return results

# Integration functions for Flask app
def setup_integrations(app, storage_path: str):
    """Setup integration components."""
    app.ai_manager = AIModelManager(storage_path)
    app.docker_handler = DockerRegistryHandler(storage_path)
    app.security_scanner = SecurityScanner()
    
    # Add Docker Registry API routes
    @app.route('/v2/')
    def docker_api_version():
        return jsonify({}), 200
    
    @app.route('/v2/<path:name>/manifests/<reference>')
    def docker_manifest(name, reference):
        try:
            manifest = app.docker_handler.handle_manifest(name, reference)
            return jsonify(manifest)
        except FileNotFoundError:
            abort(404)
    
    @app.route('/v2/<path:name>/blobs/<digest>')
    def docker_blob(name, digest):
        try:
            blob_path = app.docker_handler.handle_blob(name, digest)
            return send_file(blob_path)
        except FileNotFoundError:
            abort(404)

if __name__ == '__main__':
    # Example usage
    ai_manager = AIModelManager('./storage')
    
    # Test model detection
    test_file = './test_model.pt'
    if os.path.exists(test_file):
        metadata = ai_manager.detect_model_type(test_file)
        print("Model metadata:", json.dumps(metadata, indent=2))
        
        card = ai_manager.create_model_card(metadata, 'test_model', '1.0.0', 'Test model')
        print("Model card:", card)