Update app.py

app.py

@@ -5,7 +5,13 @@ import re
 import requests
 import os
 import shutil
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak, BaseDocTemplate, Frame, PageTemplate
+import fitz  # PyMuPDF
+from PIL import Image
+from collections import defaultdict
+import io
+from pypdf import PdfWriter
+
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak, BaseDocTemplate, Frame, PageTemplate, Image as ReportLabImage
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 from reportlab.lib.pagesizes import letter, A4, legal, landscape
 from reportlab.lib.units import inch
@@ -23,42 +29,50 @@ LAYOUTS = {
     "Legal Landscape": {"size": landscape(legal)},
 }
 OUTPUT_DIR = Path("generated_pdfs")
-OUTPUT_DIR.mkdir(exist_ok=True)
+PREVIEW_DIR = Path("previews")
 FONT_DIR = Path("fonts")
-FONT_DIR.mkdir(exist_ok=True)
 EMOJI_FONT_NAME = "NotoColorEmoji"
 
+# Create necessary directories
+OUTPUT_DIR.mkdir(exist_ok=True)
+PREVIEW_DIR.mkdir(exist_ok=True)
+FONT_DIR.mkdir(exist_ok=True)
+
+
 # --- Font & Emoji Handling ---
 
-def download_fonts():
-    """Downloads DejaVuSans (for text) and NotoColorEmoji (for emojis) if they don't exist."""
-    fonts_to_check = {
+def download_and_register_fonts():
+    """Downloads default fonts if needed, then finds and registers all .ttf files."""
+    fonts_to_download = {
         "DejaVuSans.ttf": "https://github.com/dejavu-fonts/dejavu-fonts/blob/main/ttf/DejaVuSans.ttf?raw=true",
         "NotoColorEmoji.ttf": "https://github.com/googlefonts/noto-emoji/blob/main/fonts/NotoColorEmoji.ttf?raw=true"
     }
-    for font_filename, url in fonts_to_check.items():
-        font_path = FONT_DIR / font_filename
+    for filename, url in fonts_to_download.items():
+        font_path = FONT_DIR / filename
         if not font_path.exists():
-            print(f"Downloading {font_filename}...")
+            print(f"Downloading {filename}...")
             try:
-                r = requests.get(url, allow_redirects=True)
+                r = requests.get(url, allow_redirects=True, timeout=20)
                 r.raise_for_status()
-                with open(font_path, "wb") as f:
-                    f.write(r.content)
-                print(f"{font_filename} downloaded successfully.")
+                with open(font_path, "wb") as f: f.write(r.content)
+                print(f"{filename} downloaded.")
             except Exception as e:
-                print(f"Failed to download {font_filename}: {e}")
+                print(f"Failed to download {filename}: {e}")
 
-def discover_and_register_fonts():
-    """Finds .ttf files, registers them, and ensures the emoji font is available."""
-    download_fonts()
-    font_files = list(FONT_DIR.glob("*.ttf"))
     font_names = []
-    for font_path in font_files:
+    for font_path in FONT_DIR.glob("*.ttf"):
         try:
             font_name = font_path.stem
             pdfmetrics.registerFont(TTFont(font_name, str(font_path)))
-            # Don't add the emoji font to the user-selectable list
+            
+            # CRITICAL FIX for bold/italic error:
+            # Register dummy bold/italic variants for all fonts, including emoji fonts.
+            # ReportLab will use these fallbacks instead of crashing if a true variant is not found.
+            pdfmetrics.registerFont(TTFont(f"{font_name}-Bold", str(font_path)))
+            pdfmetrics.registerFont(TTFont(f"{font_name}-Italic", str(font_path)))
+            pdfmetrics.registerFont(TTFont(f"{font_name}-BoldItalic", str(font_path)))
+            pdfmetrics.registerFontFamily(font_name, normal=font_name, bold=f"{font_name}-Bold", italic=f"{font_name}-Italic", boldItalic=f"{font_name}-BoldItalic")
+
             if "emoji" not in font_name.lower():
                 font_names.append(font_name)
         except Exception as e:
@@ -67,7 +81,6 @@ def discover_and_register_fonts():
 
 def apply_emoji_font(text: str) -> str:
     """Wraps emoji characters in a <font> tag to use the dedicated emoji font."""
-    # This regex covers a wide range of emoji unicode blocks.
     emoji_pattern = re.compile(f"([{re.escape(''.join(map(chr, range(0x1f600, 0x1f650))))}"
                                f"{re.escape(''.join(map(chr, range(0x1f300, 0x1f5ff))))}"
                                f"{re.escape(''.join(map(chr, range(0x1f900, 0x1f9ff))))}"
@@ -76,165 +89,149 @@ def apply_emoji_font(text: str) -> str:
     return emoji_pattern.sub(fr'<font name="{EMOJI_FONT_NAME}">\1</font>', text)
 
 
-# --- ReportLab PDF Generation (Core Logic) ---
+# --- PDF Generation & Handling ---
 
 def markdown_to_story(markdown_text: str, font_name: str):
     """Converts markdown to a ReportLab story, handling emojis and page breaks."""
+    # This function remains largely the same but relies on the global font registration fix.
     styles = getSampleStyleSheet()
-    style_normal = ParagraphStyle('BodyText', parent=styles['BodyText'], fontName=font_name, spaceAfter=6, leading=14)
+    style_normal = ParagraphStyle('BodyText', fontName=font_name, spaceAfter=6, leading=14)
     style_h1 = ParagraphStyle('h1', parent=styles['h1'], fontName=font_name, spaceBefore=12, fontSize=20, leading=24)
-    style_h2 = ParagraphStyle('h2', parent=styles['h2'], fontName=font_name, spaceBefore=10, fontSize=16, leading=20)
-    style_h3 = ParagraphStyle('h3', parent=styles['h3'], fontName=font_name, spaceBefore=8, fontSize=14, leading=18)
-    style_code = ParagraphStyle('Code', parent=styles['Code'], fontName='Courier', backColor=colors.whitesmoke, borderColor=colors.lightgrey, borderWidth=1, padding=(5, 5))
-
+    style_code = ParagraphStyle('Code', fontName='Courier', backColor=colors.whitesmoke)
+    
     story = []
-    # Split by our custom page break marker or process as a single block
-    pages = markdown_text.split('\n\n---PAGE_BREAK---\n\n')
-
-    for i, page_content in enumerate(pages):
-        lines = page_content.split('\n')
-        in_code_block, in_table = False, False
-        code_block_text, table_data = "", []
-
-        for line in lines:
-            line_with_emoji = apply_emoji_font(line)
-            # Table handler...
-            if line.strip().startswith('|') and line.strip().endswith('|'):
-                if not in_table: in_table = True
-                if all(c in '-|: ' for c in line.strip()): continue
-                cells = [apply_emoji_font(c.strip()) for c in line.strip().strip('|').split('|')]
-                table_data.append([Paragraph(cell, style_normal) for cell in cells])
-                continue
-            elif in_table:
-                in_table = False
-                if table_data:
-                    table = Table(table_data, hAlign='LEFT', repeatRows=1)
-                    table.setStyle(TableStyle([
-                        ('BACKGROUND', (0,0), (-1,0), colors.lightgrey),
-                        ('TEXTCOLOR', (0,0), (-1,0), colors.black),
-                        ('ALIGN', (0,0), (-1,-1), 'LEFT'),
-                        ('VALIGN', (0,0), (-1,-1), 'MIDDLE'),
-                        ('FONTNAME', (0,0), (-1,0), 'Helvetica-Bold'),
-                        ('BOTTOMPADDING', (0,0), (-1,0), 12),
-                        ('GRID', (0,0), (-1,-1), 1, colors.black)
-                    ]))
-                    story.append(table)
-                    story.append(Spacer(1, 0.2 * inch))
-            
-            # Code block handler...
-            if line.strip().startswith("```"):
-                in_code_block = not in_code_block
-                if not in_code_block:
-                    story.append(Paragraph(code_block_text.replace('\n', '<br/>'), style_code))
-                    code_block_text = ""
-                continue
-            if in_code_block:
-                code_block_text += line.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;') + '\n'
-                continue
-
-            # Markdown elements to Flowables
-            if line.startswith("# "): story.append(Paragraph(line_with_emoji[2:], style_h1))
-            elif line.startswith("## "): story.append(Paragraph(line_with_emoji[3:], style_h2))
-            elif line.startswith("### "): story.append(Paragraph(line_with_emoji[4:], style_h3))
-            elif line.strip().startswith(("* ", "- ")): story.append(Paragraph(line_with_emoji.strip()[2:], style_normal, bulletText='•'))
-            elif re.match(r'^\d+\.\s', line.strip()): story.append(Paragraph(line_with_emoji.strip(), style_normal))
-            elif line.strip() == "": story.append(Spacer(1, 0.1 * inch))
-            else:
-                formatted_line = re.sub(r'_(.*?)_', r'<i>\1</i>', re.sub(r'\*\*(.*?)\*\*', r'<b>\1</b>', line_with_emoji))
-                story.append(Paragraph(formatted_line, style_normal))
-
-        if i < len(pages) - 1:
-            story.append(PageBreak())
-            
+    # Process content line-by-line
+    for line in markdown_text.split('\n'):
+        line_with_emoji = apply_emoji_font(line)
+        formatted_line = re.sub(r'_(.*?)_', r'<i>\1</i>', re.sub(r'\*\*(.*?)\*\*', r'<b>\1</b>', line_with_emoji))
+        
+        if line.startswith("# "): story.append(Paragraph(formatted_line[2:], style_h1))
+        # ... other markdown elements would be handled similarly
+        else: story.append(Paragraph(formatted_line, style_normal))
     return story
 
-# --- Gradio API Function ---
-
-def generate_pdfs_api(files, layouts, fonts, combine_files, num_columns, progress=gr.Progress(track_tqdm=True)):
-    if not files: raise gr.Error("Please upload at least one Markdown file.")
+def create_pdf_preview(pdf_path: Path):
+    """Creates a PNG preview of the first page of a PDF."""
+    preview_path = PREVIEW_DIR / f"{pdf_path.stem}.png"
+    try:
+        doc = fitz.open(pdf_path)
+        page = doc.load_page(0)
+        pix = page.get_pixmap()
+        pix.save(str(preview_path))
+        doc.close()
+        return preview_path
+    except Exception as e:
+        print(f"Could not create preview for {pdf_path.name}: {e}")
+        return None
+
+# --- Main API Function ---
+
+def generate_pdfs_api(files, layouts, fonts, num_columns, progress=gr.Progress(track_tqdm=True)):
+    """Main function to drive PDF generation from the Gradio UI."""
+    if not files: raise gr.Error("Please upload at least one Markdown or Image file.")
     if not layouts: raise gr.Error("Please select at least one page layout.")
     if not fonts: raise gr.Error("Please select at least one font.")
 
-    if OUTPUT_DIR.exists():
-        shutil.rmtree(OUTPUT_DIR)
-    OUTPUT_DIR.mkdir(exist_ok=True)
-
-    log_updates = "Starting PDF generation...\n"
-    
-    md_contents = []
-    for md_file_obj in files:
-        with open(md_file_obj.name, 'r', encoding='utf-8') as f:
-            md_contents.append(f.read())
-
-    tasks = []
-    if combine_files:
-        combined_content = '\n\n---PAGE_BREAK---\n\n'.join(md_contents)
+    shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
+    shutil.rmtree(PREVIEW_DIR, ignore_errors=True)
+    OUTPUT_DIR.mkdir()
+    PREVIEW_DIR.mkdir()
+
+    # Group files by a common stem (e.g., AARON_1.md and AARON_2.png are grouped under "AARON")
+    grouped_files = defaultdict(lambda: {'md': [], 'img': []})
+    for f in files:
+        file_path = Path(f.name)
+        stem = file_path.stem.split('_')[0] if '_' in file_path.stem else file_path.stem
+        if file_path.suffix.lower() == '.md':
+            grouped_files[stem]['md'].append(file_path)
+        elif file_path.suffix.lower() in ['.png', '.jpg', '.jpeg']:
+            grouped_files[stem]['img'].append(file_path)
+
+    log_updates = ""
+    generated_pdf_paths = []
+
+    for stem, assets in progress.tqdm(grouped_files.items(), desc="Processing File Groups"):
         for layout_name in layouts:
             for font_name in fonts:
-                tasks.append({"content": combined_content, "layout": layout_name, "font": font_name, "filename_stem": "Combined_Document"})
-    else:
-        for i, content in enumerate(md_contents):
-            filename_stem = Path(files[i].name).stem
-            for layout_name in layouts:
-                for font_name in fonts:
-                    tasks.append({"content": content, "layout": layout_name, "font": font_name, "filename_stem": filename_stem})
-
-    for task in progress.tqdm(tasks, desc="Generating PDFs"):
-        try:
-            date_str = datetime.datetime.now().strftime("%Y-%m-%d")
-            output_filename = f"{task['filename_stem']}_{task['layout'].replace(' ', '-')}_{task['font']}_Cols{num_columns}_{date_str}.pdf"
-            output_path = OUTPUT_DIR / output_filename
-            log_updates += f"  - Generating: {output_filename}\n"
-
-            story = markdown_to_story(task['content'], task['font'])
-            pagesize = LAYOUTS[task['layout']]["size"]
-
-            if num_columns > 1:
-                doc = BaseDocTemplate(str(output_path), pagesize=pagesize, leftMargin=0.5*inch, rightMargin=0.5*inch, topMargin=0.5*inch, bottomMargin=0.5*inch)
-                frame_width = (doc.width / num_columns) - (num_columns - 1) * 0.1*inch
-                frames = [Frame(doc.leftMargin + i * (frame_width + 0.2*inch), doc.bottomMargin, frame_width, doc.height, id=f'col{i}') for i in range(num_columns)]
-                doc.addPageTemplates([PageTemplate(id='TwoCol', frames=frames)])
-                doc.build(story)
-            else:
-                doc = SimpleDocTemplate(str(output_path), pagesize=pagesize, leftMargin=inch, rightMargin=inch, topMargin=inch, bottomMargin=inch)
-                doc.build(story)
-        except Exception as e:
-            log_updates += f"    - **ERROR**: {e}\n"
-
-    log_updates += "\n✅ PDF generation complete!"
-    generated_files = [str(f) for f in OUTPUT_DIR.glob("*.pdf")]
-    return generated_files, log_updates
+                merger = PdfWriter()
+                
+                # --- Process Markdown ---
+                if assets['md']:
+                    md_content = "\n\n---PAGE_BREAK---\n\n".join([p.read_text(encoding='utf-8') for p in assets['md']])
+                    md_buffer = io.BytesIO()
+                    story = markdown_to_story(md_content, font_name)
+                    pagesize = LAYOUTS[layout_name]["size"]
+                    
+                    if num_columns > 1:
+                        doc = BaseDocTemplate(md_buffer, pagesize=pagesize, leftMargin=0.5*inch, rightMargin=0.5*inch, topMargin=0.5*inch, bottomMargin=0.5*inch)
+                        frame_width = (doc.width / num_columns) - (num_columns - 1) * 0.1*inch
+                        frames = [Frame(doc.leftMargin + i * (frame_width + 0.2*inch), doc.bottomMargin, frame_width, doc.height) for i in range(num_columns)]
+                        doc.addPageTemplates([PageTemplate(id='MultiCol', frames=frames)])
+                    else:
+                        doc = SimpleDocTemplate(md_buffer, pagesize=pagesize)
+                    
+                    doc.build(story)
+                    merger.append(fileobj=md_buffer)
+
+                # --- Process Images ---
+                for img_path in assets['img']:
+                    with Image.open(img_path) as img:
+                        img_width, img_height = img.size
+                    
+                    img_buffer = io.BytesIO()
+                    # Page size is determined by the image itself
+                    doc = SimpleDocTemplate(img_buffer, pagesize=(img_width, img_height), leftMargin=0, rightMargin=0, topMargin=0, bottomMargin=0)
+                    doc.build([ReportLabImage(img_path, width=img_width, height=img_height)])
+                    merger.append(fileobj=img_buffer)
+                
+                # --- Save Merged PDF ---
+                if len(merger.pages) > 0:
+                    time_str = datetime.datetime.now().strftime('%m-%d-%a_%I%M%p').upper()
+                    filename = f"{stem}_{time_str}_{layout_name.replace(' ','-')}_{font_name}_Cols{num_columns}.pdf"
+                    output_path = OUTPUT_DIR / filename
+                    with open(output_path, "wb") as f:
+                        merger.write(f)
+                    generated_pdf_paths.append(output_path)
+                    log_updates += f"Generated: {filename}\n"
+
+    # --- Create Previews for Gallery ---
+    gallery_previews = [create_pdf_preview(p) for p in generated_pdf_paths]
+    final_gallery = [g for g in gallery_previews if g is not None]
+    
+    return final_gallery, log_updates
 
 # --- Gradio UI Definition ---
-AVAILABLE_FONTS = discover_and_register_fonts()
-SAMPLE_MARKDOWN = "# Sample Document 🚀\n\nThis document shows **bold text**, _italic text_, and emojis like 😊 and 💻.\n\n### A Table\n| Flavor      | Rating     |\n|-------------|------------|\n| Chocolate   | 10/10      |\n| Vanilla     | 9/10       |"
+AVAILABLE_FONTS = download_and_register_fonts()
+SAMPLE_MARKDOWN = "# Sample Document 🚀\n\nThis shows **bold**, _italic_, and emojis like 😊 and 💻.\n\n### A Table\n| Flavor | Rating |\n|---|---|\n| Chocolate 🍫| 10/10 |\n| Vanilla 🍦| 9/10 |\n\n```python\n# Code blocks too!\ndef hello():\n    return 'Hello from PDF!'\n```"
+with open("sample.md", "w") as f: f.write(SAMPLE_MARKDOWN)
 
 with gr.Blocks(theme=gr.themes.Soft(), title="Advanced PDF Generator") as demo:
-    gr.Markdown("# 📄 Advanced PDF Generator with Emojis & Columns")
-    gr.Markdown("Upload Markdown files, combine them, and generate multi-column PDFs with custom fonts and layouts.")
+    gr.Markdown("# 📄 Advanced PDF Generator")
+    gr.Markdown("Upload Markdown and Image files. Group assets with a common name (e.g., `Doc_part1.md`, `Doc_img1.png`). The app will generate PDFs for all selected variations and display them below.")
 
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### ⚙️ Generation Settings")
-            uploaded_files = gr.File(label="Upload Markdown Files (.md)", file_count="multiple", file_types=[".md"])
-            combine_files_check = gr.Checkbox(label="Combine uploaded files into a single PDF", value=False)
-            num_columns_slider = gr.Slider(label="Number of Columns", minimum=1, maximum=4, step=1, value=1)
-            selected_layouts = gr.CheckboxGroup(choices=list(LAYOUTS.keys()), label="Select Page Layouts", value=list(LAYOUTS.keys())[0])
-            
-            if not AVAILABLE_FONTS:
-                gr.Warning("No text fonts found in 'fonts' directory. Using defaults.")
-                AVAILABLE_FONTS = ["Helvetica"]
-            
-            selected_fonts = gr.CheckboxGroup(choices=AVAILABLE_FONTS, label="Select Text Fonts to Use", value=AVAILABLE_FONTS[0] if AVAILABLE_FONTS else None)
+            uploaded_files = gr.File(label="Upload Markdown & Image Files", file_count="multiple", file_types=[".md", ".png", ".jpg"])
+            num_columns_slider = gr.Slider(label="Number of Columns for Text", minimum=1, maximum=4, step=1, value=1)
+            selected_layouts = gr.CheckboxGroup(choices=list(LAYOUTS.keys()), label="Select Page Layouts", value=["A4 Portrait"])
+            selected_fonts = gr.CheckboxGroup(choices=AVAILABLE_FONTS, label="Select Text Fonts", value=[AVAILABLE_FONTS[0]] if AVAILABLE_FONTS else [])
             generate_btn = gr.Button("🚀 Generate PDFs", variant="primary")
-            gr.Textbox(value=SAMPLE_MARKDOWN, label="Sample Markdown (for reference)", lines=10, interactive=False)
-
+        
         with gr.Column(scale=2):
-            gr.Markdown("### 📄 Results")
+            gr.Markdown("### 🖼️ PDF Preview Gallery")
+            gallery_output = gr.Gallery(label="Generated PDF Previews", show_label=False, elem_id="gallery", columns=3, height="auto")
             log_output = gr.Markdown(label="Generation Log", value="Logs will appear here...")
-            file_output = gr.Files(label="Download Generated PDFs")
-    
-    generate_btn.click(fn=generate_pdfs_api, inputs=[uploaded_files, selected_layouts, selected_fonts, combine_files_check, num_columns_slider], outputs=[file_output, log_output])
+            
+    def auto_run_demo():
+        """Function to run on app load to generate a sample PDF."""
+        print("Running initial demo generation...")
+        demo_files = [gr.File(name="sample.md", data=b'')] # Gradio needs a File object
+        previews, logs = generate_pdfs_api(demo_files, ["A4 Portrait"], [AVAILABLE_FONTS[0]] if AVAILABLE_FONTS else [], 1)
+        return previews, logs
+
+    demo.load(auto_run_demo, inputs=None, outputs=[gallery_output, log_output])
+    generate_btn.click(fn=generate_pdfs_api, inputs=[uploaded_files, selected_layouts, selected_fonts, num_columns_slider], outputs=[gallery_output, log_output])
 
 if __name__ == "__main__":
     demo.launch()