parallel_render.py

"""
Parallel Render Module for KiRender
Handles multi-core frame rendering using concurrent.futures
"""

import os
import sys
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor, as_completed


def get_cpu_count():
    """Get available CPU cores."""
    try:
        return os.cpu_count() or 4
    except:
        return 4


def get_recommended_workers():
    """
    Get recommended number of parallel workers for KiCad 3D rendering.
    
    KiCad's raytracing (--quality high) is CPU-bound software rendering.
    More workers = faster, but each needs ~1-2GB RAM.
    Leave 1-2 cores for system responsiveness.
    """
    cores = get_cpu_count()
    # Use most cores, leave 2 for system (or 1 if only 4 cores)
    if cores <= 4:
        return max(1, cores - 1)
    else:
        return max(1, cores - 2)


def render_single_frame(args):
    """
    Render a single frame. This function is called by parallel workers.
    
    Args:
        args: tuple of (frame_index, frame_path, cmd, frame_params)
    
    Returns:
        tuple of (frame_index, success, elapsed_time, error_message)
    """
    frame_index, frame_path, cmd = args
    
    start_time = time.time()
    
    try:
        # Suppress console window on Windows
        startupinfo = None
        creationflags = 0
        if sys.platform == "win32":
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            creationflags = subprocess.CREATE_NO_WINDOW
        
        # Debug: print command being run
        # print(f"[Frame {frame_index}] Running: {' '.join(cmd)}")
        
        result = subprocess.run(
            cmd,
            startupinfo=startupinfo,
            creationflags=creationflags,
            capture_output=True,
            text=True,
            timeout=900  # 15 minute timeout per frame (raytracing can be slow)
        )
        
        elapsed = time.time() - start_time
        
        if result.returncode != 0:
            error_msg = result.stderr.strip() if result.stderr else f"Exit code {result.returncode}"
            return (frame_index, False, elapsed, error_msg)
        
        # Verify output file exists
        if not os.path.exists(frame_path):
            return (frame_index, False, elapsed, f"Output file not created: {frame_path}")
        
        return (frame_index, True, elapsed, None)
        
    except subprocess.TimeoutExpired:
        elapsed = time.time() - start_time
        return (frame_index, False, elapsed, f"Timeout: frame took longer than 15 minutes")
    except Exception as e:
        elapsed = time.time() - start_time
        return (frame_index, False, elapsed, f"Exception: {str(e)}")
        return (frame_index, False, elapsed, str(e))


class ParallelFrameRenderer:
    """
    Renders multiple frames in parallel using a thread pool.
    
    Usage:
        renderer = ParallelFrameRenderer(num_workers=4)
        results = renderer.render_frames(frame_tasks, progress_callback)
    """
    
    def __init__(self, num_workers=None):
        """
        Initialize parallel renderer.
        
        Args:
            num_workers: Number of parallel workers. None = auto (cpu_count - 1)
        """
        if num_workers is None or num_workers <= 0:
            self.num_workers = get_recommended_workers()
        else:
            self.num_workers = min(num_workers, get_cpu_count())
        
        self.cancelled = False
        self.completed_frames = 0
        self.total_frames = 0
        self.frame_times = []
    
    def cancel(self):
        """Signal cancellation to stop rendering."""
        self.cancelled = True
    
    def render_frames(self, frame_tasks, progress_callback=None, log_callback=None):
        """
        Render all frames in parallel.
        
        Args:
            frame_tasks: List of (frame_index, frame_path, cmd) tuples
            progress_callback: Optional callback(completed, total, avg_time, eta_str)
            log_callback: Optional callback(message) for logging
        
        Returns:
            tuple of (success_count, fail_count, avg_frame_time, errors)
        """
        self.total_frames = len(frame_tasks)
        self.completed_frames = 0
        self.frame_times = []
        self.cancelled = False
        
        errors = []
        success_count = 0
        fail_count = 0
        
        if log_callback:
            log_callback(f"🚀 Starting parallel render with {self.num_workers} workers...")
        
        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
            # Submit all tasks
            future_to_frame = {
                executor.submit(render_single_frame, task): task[0]
                for task in frame_tasks
            }
            
            # Process completed frames as they finish
            for future in as_completed(future_to_frame):
                if self.cancelled:
                    # Cancel remaining futures
                    for f in future_to_frame:
                        f.cancel()
                    if log_callback:
                        log_callback("Render cancelled by user.")
                    break
                
                frame_index = future_to_frame[future]
                
                try:
                    idx, success, elapsed, error = future.result()
                    
                    self.frame_times.append(elapsed)
                    self.completed_frames += 1
                    
                    if success:
                        success_count += 1
                    else:
                        fail_count += 1
                        errors.append((idx, error))
                        if log_callback:
                            log_callback(f"⚠️ Frame {idx} error: {error}")
                    
                    # Calculate ETA
                    if self.frame_times:
                        avg_time = sum(self.frame_times) / len(self.frame_times)
                        # Account for parallel workers in ETA
                        frames_remaining = self.total_frames - self.completed_frames
                        batches_remaining = frames_remaining / self.num_workers
                        eta_seconds = avg_time * batches_remaining
                        eta_str = self._format_eta(eta_seconds)
                        
                        if progress_callback:
                            progress_callback(
                                self.completed_frames,
                                self.total_frames,
                                avg_time,
                                eta_str
                            )
                
                except Exception as e:
                    fail_count += 1
                    errors.append((frame_index, str(e)))
        
        avg_frame_time = sum(self.frame_times) / len(self.frame_times) if self.frame_times else 0
        
        return (success_count, fail_count, avg_frame_time, errors)
    
    def _format_eta(self, seconds):
        """Format seconds into human-readable ETA."""
        if seconds < 60:
            return f"{seconds:.0f}s"
        elif seconds < 3600:
            mins = int(seconds // 60)
            secs = int(seconds % 60)
            return f"{mins}m {secs}s"
        else:
            hrs = int(seconds // 3600)
            mins = int((seconds % 3600) // 60)
            return f"{hrs}h {mins}m"


def build_frame_tasks(params, temp_dir):
    """
    Build list of frame rendering tasks.
    
    Args:
        params: Render parameters dict
        temp_dir: Temporary directory for frame output
    
    Returns:
        List of (frame_index, frame_path, cmd) tuples
    """
    from .cli_builder import build_cli_command
    
    tasks = []
    step_angle = (360.0 * params['revolutions']) / params['frames']
    
    for i in range(params['frames']):
        angle = i * step_angle
        
        # Build rotation string based on axis
        if params['axis'] == 0:  # Y Axis (Spin)
            rot_str = f"{params['tilt_x']},{angle},{params['tilt_z']}"
        elif params['axis'] == 1:  # Z Axis (Table/Turntable)
            rot_str = f"{params['tilt_x']},{params['tilt_y']},{angle}"
        else:  # X Axis (Flip)
            rot_str = f"{angle},{params['tilt_y']},{params['tilt_z']}"
        
        frame_name = f"frame_{i:04d}.png"
        frame_path = os.path.join(temp_dir, frame_name)
        
        # Build CLI params for this frame
        frame_params = {
            'width': params['width'],
            'height': params['height'],
            'zoom': params['zoom'],
            'background': params['background'],
            'quality': params['quality'],
            'rotate': rot_str,
            'side': params.get('side', 'top'),
            'perspective': params.get('perspective') != "orthographic",
        }
        
        # Add lighting parameters if specified
        for light_key in ('light_top', 'light_bottom', 'light_side',
                          'light_camera', 'light_side_elevation'):
            if light_key in params:
                frame_params[light_key] = params[light_key]
        
        cmd = build_cli_command(frame_params, params['kicad_cli'], params['pcb_path'], frame_path)
        
        tasks.append((i, frame_path, cmd))
    
    return tasks