ETPlot/orchestration/pdf_export.py

"""
PDF Export Module for Scientific Gallery Generator

This module handles exporting and merging multiple plots into a single PDF
using PyPDF2 and reportlab for layout.
"""

import io
import tempfile
from pathlib import Path
from typing import List, Dict, Any
import subprocess

try:
    from PyPDF2 import PdfReader, PdfWriter
    from reportlab.pdfgen import canvas
    from reportlab.lib.pagesizes import letter, A4
    from reportlab.lib.utils import ImageReader
    HAS_PDF_LIBS = True
except ImportError:
    HAS_PDF_LIBS = False


class PDFExporter:
    """Handles exporting multiple plots to a merged PDF"""
    
    def __init__(self):
        self.page_size = A4
        self.margin = 50
        
    def merge_plots(self, plot_paths: List[str], layout: Dict[str, int], 
                   output_path: str = None) -> bytes:
        """
        Merge multiple PDF plots into a single PDF with grid layout.
        
        Args:
            plot_paths: List of paths to PDF files
            layout: Dictionary with 'rows' and 'cols' keys
            output_path: Optional output file path
            
        Returns:
            PDF bytes
        """
        if not HAS_PDF_LIBS:
            return self._merge_with_pdfjam(plot_paths, layout, output_path)
        
        return self._merge_with_pypdf(plot_paths, layout, output_path)
    
    def _merge_with_pypdf(self, plot_paths: List[str], layout: Dict[str, int], 
                         output_path: str = None) -> bytes:
        """Merge PDFs using PyPDF2 and reportlab"""
        from reportlab.pdfgen import canvas
        from reportlab.lib.pagesizes import A4
        
        # Create a new PDF with the layout
        buffer = io.BytesIO()
        c = canvas.Canvas(buffer, pagesize=A4)
        page_width, page_height = A4
        
        rows = layout['rows']
        cols = layout['cols']
        
        # Calculate dimensions for each plot
        plot_width = (page_width - 2 * self.margin) / cols
        plot_height = (page_height - 2 * self.margin) / rows
        
        # Place each plot in the grid
        for i, plot_path in enumerate(plot_paths):
            if i >= rows * cols:
                break
                
            row = i // cols
            col = i % cols
            
            # Calculate position
            x = self.margin + col * plot_width
            y = page_height - self.margin - (row + 1) * plot_height
            
            try:
                # Read the source PDF
                with open(plot_path, 'rb') as f:
                    reader = PdfReader(f)
                    if len(reader.pages) > 0:
                        page = reader.pages[0]
                        
                        # Convert PDF page to image and place it
                        # This is a simplified approach - in practice you'd want
                        # to properly scale and position the PDF content
                        self._draw_pdf_placeholder(c, x, y, plot_width, plot_height, 
                                                 Path(plot_path).stem)
                        
            except Exception as e:
                print(f"Error processing {plot_path}: {e}")
                self._draw_error_placeholder(c, x, y, plot_width, plot_height)
        
        c.save()
        
        if output_path:
            with open(output_path, 'wb') as f:
                f.write(buffer.getvalue())
        
        return buffer.getvalue()
    
    def _merge_with_pdfjam(self, plot_paths: List[str], layout: Dict[str, int], 
                          output_path: str = None) -> bytes:
        """Merge PDFs using pdfjam (requires pdfpages LaTeX package)"""
        rows = layout['rows']
        cols = layout['cols']
        
        # Create temporary output file
        with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as tmp_file:
            temp_output = tmp_file.name
        
        try:
            # Build pdfjam command
            cmd = [
                'pdfjam',
                '--nup', f'{cols}x{rows}',
                '--landscape' if cols > rows else '--no-landscape',
                '--frame', 'true',
                '--delta', '10pt 10pt',
                '--offset', '0pt 0pt',
                '--outfile', temp_output
            ]
            
            # Add input files
            cmd.extend(plot_paths[:rows * cols])
            
            # Run pdfjam
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode != 0:
                raise Exception(f"pdfjam failed: {result.stderr}")
            
            # Read the output file
            with open(temp_output, 'rb') as f:
                pdf_bytes = f.read()
            
            if output_path:
                with open(output_path, 'wb') as f:
                    f.write(pdf_bytes)
            
            return pdf_bytes
            
        finally:
            # Clean up temporary file
            Path(temp_output).unlink(missing_ok=True)
    
    def _draw_pdf_placeholder(self, canvas, x: float, y: float, width: float, 
                             height: float, plot_name: str):
        """Draw a placeholder for a PDF plot"""
        # Draw border
        canvas.setStrokeColorRGB(0.5, 0.5, 0.5)
        canvas.setLineWidth(1)
        canvas.rect(x, y, width, height)
        
        # Draw plot name
        canvas.setFillColorRGB(0, 0, 0)
        canvas.setFont("Helvetica", 10)
        text_width = canvas.stringWidth(plot_name, "Helvetica", 10)
        text_x = x + (width - text_width) / 2
        text_y = y + height / 2
        canvas.drawString(text_x, text_y, plot_name)
    
    def _draw_error_placeholder(self, canvas, x: float, y: float, width: float, 
                               height: float):
        """Draw an error placeholder"""
        # Draw red border
        canvas.setStrokeColorRGB(1, 0, 0)
        canvas.setLineWidth(2)
        canvas.rect(x, y, width, height)
        
        # Draw error text
        canvas.setFillColorRGB(1, 0, 0)
        canvas.setFont("Helvetica-Bold", 12)
        error_text = "Error loading plot"
        text_width = canvas.stringWidth(error_text, "Helvetica-Bold", 12)
        text_x = x + (width - text_width) / 2
        text_y = y + height / 2
        canvas.drawString(text_x, text_y, error_text)


def check_dependencies() -> Dict[str, bool]:
    """Check if required dependencies are available"""
    deps = {
        'pypdf2': HAS_PDF_LIBS,
        'pdfjam': False
    }
    
    # Check for pdfjam
    try:
        result = subprocess.run(['pdfjam', '--version'], 
                              capture_output=True, text=True)
        deps['pdfjam'] = result.returncode == 0
    except FileNotFoundError:
        pass
    
    return deps


if __name__ == "__main__":
    # Test the exporter
    exporter = PDFExporter()
    deps = check_dependencies()
    print("Available dependencies:", deps)