ritz/module_system.py

"""
Ritz Module System - Progressive Loading Concept
Split apps into multiple WASM modules for security and performance
"""

import hashlib
import json
from dataclasses import dataclass
from typing import Dict, List, Optional, Set

@dataclass
class WASMModule:
    """Represents a WASM module"""
    name: str
    path: str
    size_kb: int
    dependencies: List[str]
    exports: List[str]
    requires_auth: bool = False
    cache_strategy: str = 'aggressive'  # aggressive, normal, no-cache

class ModuleGraph:
    """Dependency graph for WASM modules"""
    
    def __init__(self):
        self.modules: Dict[str, WASMModule] = {}
        self.loaded: Set[str] = set()
    
    def add_module(self, module: WASMModule):
        self.modules[module.name] = module
    
    def get_load_order(self, module_name: str) -> List[str]:
        """Get modules to load in dependency order"""
        if module_name in self.loaded:
            return []
        
        order = []
        visited = set()
        
        def visit(name):
            if name in visited or name in self.loaded:
                return
            visited.add(name)
            
            module = self.modules.get(name)
            if module:
                for dep in module.dependencies:
                    visit(dep)
                order.append(name)
        
        visit(module_name)
        return order

class RitzModuleSystem:
    """Progressive module loading system"""
    
    def __init__(self):
        self.graph = ModuleGraph()
        self.setup_modules()
    
    def setup_modules(self):
        """Define application module structure"""
        
        # Bootstrap - minimal, always loaded
        self.graph.add_module(WASMModule(
            name='bootstrap',
            path='/wasm/bootstrap.wasm',
            size_kb=50,
            dependencies=[],
            exports=['LoginForm', 'ModuleLoader', 'AuthService'],
            requires_auth=False,
            cache_strategy='aggressive'
        ))
        
        # Core app - loaded after authentication
        self.graph.add_module(WASMModule(
            name='app-core',
            path='/wasm/app-core.wasm',
            size_kb=200,
            dependencies=['bootstrap'],
            exports=['Router', 'Dashboard', 'Navigation'],
            requires_auth=True,
            cache_strategy='aggressive'
        ))
        
        # Feature modules - lazy loaded
        self.graph.add_module(WASMModule(
            name='admin',
            path='/wasm/features/admin.wasm',
            size_kb=150,
            dependencies=['app-core'],
            exports=['AdminPanel', 'UserManagement', 'Settings'],
            requires_auth=True,
            cache_strategy='normal'
        ))
        
        self.graph.add_module(WASMModule(
            name='analytics',
            path='/wasm/features/analytics.wasm',
            size_kb=120,
            dependencies=['app-core', 'charts'],
            exports=['AnalyticsDashboard', 'Reports'],
            requires_auth=True,
            cache_strategy='normal'
        ))
        
        self.graph.add_module(WASMModule(
            name='charts',
            path='/wasm/lib/charts.wasm',
            size_kb=80,
            dependencies=[],
            exports=['LineChart', 'BarChart', 'PieChart'],
            requires_auth=True,
            cache_strategy='aggressive'
        ))
    
    def generate_loader_js(self) -> str:
        """Generate JavaScript module loader"""
        return '''
class RitzModuleLoader {
    constructor() {
        this.modules = new Map();
        this.loading = new Map();
        this.cache = caches.open('ritz-modules-v1');
    }
    
    async loadModule(name, config = {}) {
        // Check if already loaded
        if (this.modules.has(name)) {
            return this.modules.get(name);
        }
        
        // Check if already loading (prevent double-load)
        if (this.loading.has(name)) {
            return this.loading.get(name);
        }
        
        // Start loading
        const loadPromise = this._loadModule(name, config);
        this.loading.set(name, loadPromise);
        
        try {
            const module = await loadPromise;
            this.modules.set(name, module);
            return module;
        } finally {
            this.loading.delete(name);
        }
    }
    
    async _loadModule(name, config) {
        const manifest = await this.getManifest();
        const moduleInfo = manifest[name];
        
        if (!moduleInfo) {
            throw new Error(`Module ${name} not found`);
        }
        
        // Check authentication requirement
        if (moduleInfo.requires_auth && !this.isAuthenticated()) {
            throw new Error(`Module ${name} requires authentication`);
        }
        
        // Load dependencies first
        for (const dep of moduleInfo.dependencies) {
            await this.loadModule(dep, config);
        }
        
        // Try cache first
        const cache = await this.cache;
        const cached = await cache.match(moduleInfo.path);
        
        let response;
        if (cached && !config.bypassCache) {
            console.log(`Loading ${name} from cache`);
            response = cached;
        } else {
            console.log(`Downloading ${name} (${moduleInfo.size_kb}KB)`);
            
            // Add auth token if required
            const headers = {};
            if (moduleInfo.requires_auth) {
                headers['Authorization'] = `Bearer ${this.getAuthToken()}`;
            }
            
            response = await fetch(moduleInfo.path, { headers });
            
            // Cache based on strategy
            if (moduleInfo.cache_strategy !== 'no-cache') {
                await cache.put(moduleInfo.path, response.clone());
            }
        }
        
        // Instantiate WASM
        const bytes = await response.arrayBuffer();
        const wasm = await WebAssembly.instantiate(bytes, {
            env: this.getImports(name),
            ...this.getModuleImports(moduleInfo.dependencies)
        });
        
        return wasm.instance;
    }
    
    getModuleImports(dependencies) {
        // Link imported functions from dependencies
        const imports = {};
        for (const dep of dependencies) {
            const module = this.modules.get(dep);
            if (module) {
                imports[dep] = module.exports;
            }
        }
        return imports;
    }
    
    async preloadModules(names) {
        // Preload modules in parallel
        const promises = names.map(name => this.loadModule(name));
        await Promise.all(promises);
    }
    
    isAuthenticated() {
        return localStorage.getItem('auth_token') !== null;
    }
    
    getAuthToken() {
        return localStorage.getItem('auth_token');
    }
    
    async getManifest() {
        // Module manifest with metadata
        const response = await fetch('/wasm/manifest.json');
        return response.json();
    }
}

// Usage
const loader = new RitzModuleLoader();

// Bootstrap phase - no auth required
await loader.loadModule('bootstrap');

// After login
async function onLoginSuccess(token) {
    localStorage.setItem('auth_token', token);
    
    // Now load the main app
    await loader.loadModule('app-core');
    
    // Preload common features
    loader.preloadModules(['charts']);  // Non-blocking
    
    // Navigate to dashboard
    window.RitzApp.navigate('/dashboard');
}

// Lazy load on route change
window.addEventListener('route-change', async (e) => {
    const route = e.detail.route;
    
    if (route.startsWith('/admin')) {
        await loader.loadModule('admin');
    } else if (route.startsWith('/analytics')) {
        await loader.loadModule('analytics');
    }
});
        '''
    
    def generate_manifest(self) -> Dict:
        """Generate module manifest"""
        manifest = {}
        for name, module in self.graph.modules.items():
            manifest[name] = {
                'path': module.path,
                'size_kb': module.size_kb,
                'dependencies': module.dependencies,
                'requires_auth': module.requires_auth,
                'cache_strategy': module.cache_strategy,
                'hash': self.calculate_hash(module)
            }
        return manifest
    
    def calculate_hash(self, module: WASMModule) -> str:
        """Calculate hash for cache invalidation"""
        content = f"{module.name}:{module.path}:{module.size_kb}"
        return hashlib.sha256(content.encode()).hexdigest()[:12]
    
    def split_components_into_modules(self, components: List[str]) -> Dict[str, List[str]]:
        """Automatically split components into modules"""
        
        # Example heuristic: group by feature area
        modules = {
            'bootstrap': [],
            'app-core': [],
            'admin': [],
            'analytics': [],
            'reports': []
        }
        
        for component in components:
            if component.startswith('Login') or component.startswith('Auth'):
                modules['bootstrap'].append(component)
            elif component.startswith('Admin'):
                modules['admin'].append(component)
            elif component.startswith('Analytics') or component.startswith('Chart'):
                modules['analytics'].append(component)
            elif component.startswith('Report'):
                modules['reports'].append(component)
            else:
                modules['app-core'].append(component)
        
        # Remove empty modules
        return {k: v for k, v in modules.items() if v}

class SecurityExample:
    """Security benefits of modular loading"""
    
    @staticmethod
    def bootstrap_only_code():
        """Code that ships to all users (even unauthenticated)"""
        return '''
# bootstrap.py - Compiles to bootstrap.wasm
@component
def LoginPage():
    username = signal('')
    password = signal('')
    error = signal(None)
    
    async def handle_login():
        try:
            response = await fetch('/api/login', {
                'method': 'POST',
                'body': json.dumps({
                    'username': username.value,
                    'password': password.value
                })
            })
            
            if response.ok:
                token = await response.json()
                # Now load the real app
                await ModuleLoader.load('app-core')
                navigate('/dashboard')
            else:
                error.value = 'Invalid credentials'
        except Exception as e:
            error.value = str(e)
    
    with Main():
        with Form(onsubmit=handle_login):
            yield H1('Login')
            yield Input(
                type='text',
                placeholder='Username',
                value=username,
                oninput=lambda e: username.set(e.target.value)
            )
            yield Input(
                type='password',
                placeholder='Password',
                value=password,
                oninput=lambda e: password.set(e.target.value)
            )
            yield Button('Sign In', type='submit')
            
            if error.value:
                yield ErrorMessage(error.value)

# IMPORTANT: No business logic, no API endpoints, no sensitive data!
# This is ALL that unauthenticated users can see
        '''
    
    @staticmethod
    def protected_module_code():
        """Code that only authenticated users receive"""
        return '''
# app-core.py - Compiles to app-core.wasm
# This is NEVER sent to unauthenticated users

@component
def Dashboard():
    # Business logic here
    user_data = signal(None)
    
    async def load_user_data():
        # API endpoints only known to authenticated users
        response = await fetch('/api/user/profile')
        user_data.value = await response.json()
    
    # ... rest of application
        '''

# Demonstration
if __name__ == '__main__':
    print("=== Ritz Module System ===\n")
    
    system = RitzModuleSystem()
    
    print("1. Module Structure:")
    for name, module in system.graph.modules.items():
        print(f"  {name}:")
        print(f"    Size: {module.size_kb}KB")
        print(f"    Auth Required: {module.requires_auth}")
        print(f"    Dependencies: {module.dependencies}")
        print(f"    Exports: {module.exports[:2]}...")
    
    print("\n2. Load Order for 'analytics':")
    order = system.graph.get_load_order('analytics')
    print(f"  Must load in order: {order}")
    total_size = sum(system.graph.modules[m].size_kb for m in order)
    print(f"  Total download: {total_size}KB")
    
    print("\n3. Module Manifest:")
    manifest = system.generate_manifest()
    print(json.dumps(manifest, indent=2)[:400] + "...")
    
    print("\n4. Security Benefits:")
    print("  - Unauthenticated users only get bootstrap.wasm (50KB)")
    print("  - No API endpoints or business logic exposed")
    print("  - Feature modules loaded on-demand")
    print("  - Each module can have different cache strategies")
    
    print("\n5. Performance Benefits:")
    print("  - Initial load: 50KB instead of 600KB+")
    print("  - Aggressive caching for stable modules")
    print("  - Parallel preloading of likely-needed modules")
    print("  - Route-based code splitting")