#!/usr/bin/env python3
"""Generate SDS PDF with proper page numbering via two-pass process."""

from weasyprint import HTML
import pikepdf
import io

def add_page_numbers(pdf_path):
    """Post-process a PDF to add 'Page X of Y' at bottom center of each page."""
    src = pikepdf.Pdf.open(pdf_path)
    total = len(src.pages)
    
    if total <= 1:
        src.save(pdf_path)
        return total
    
    for i, page in enumerate(src.pages, 1):
        text = f"Page {i} of {total}"
        mediabox = page.mediabox
        w = float(mediabox[2])
        bottom = 50  # ~1.4cm from bottom in points
        
        # Build raw PDF content stream for page number text
        stream_data = (
            f"BT\n"
            f"/Helvetica 7 Tf\n"
            f"{w/2:.0f} {bottom} Td\n"
            f"0.357 0.400 0.435 rg\n"
            f"({text}) Tj\n"
            f"ET\n"
        ).encode('latin-1')
        
        # Add as new content stream to page
        new_stream = pikepdf.Stream(src, stream_data)
        page.contents_add(new_stream)
    
    src.save(pdf_path)
    return total


def generate_sds(html, output_path):
    """Generate SDS PDF with page numbering."""
    buf = io.BytesIO()
    HTML(string=html).write_pdf(buf)
    buf.seek(0)
    
    with open(output_path, 'wb') as f:
        f.write(buf.read())
    
    n = add_page_numbers(output_path)
    return n


if __name__ == '__main__':
    # Test with 3 pages
    test = """<html><body>
    <style>
    @page { size: A4; margin: 2cm; }
    body { font-family: sans-serif; font-size: 12pt; }
    </style>
    <p>Content page 1</p>
    <p>Content page 2</p>
    <p>Content page 3</p>
    </body></html>"""
    
    n = generate_sds(test, '/tmp/test_numbered.pdf')
    print(f'Generated {n} pages')
    
    # Verify
    with pikepdf.Pdf.open('/tmp/test_numbered.pdf') as pdf:
        for i, page in enumerate(pdf.pages, 1):
            content_streams = list(page.contents)
            print(f'Page {i}: {len(content_streams)} content stream(s)')