#!/usr/bin/env python3
import argparse
import os
import sys

def wrap_text(text, font, max_width):
    """Wrap text to fit within max_width using the given font."""
    lines = []
    # Split by explicit newlines first
    for paragraph in text.split('\n'):
        if not paragraph.strip():
            lines.append("")
            continue
        
        words = paragraph.split(' ')
        current_line = []
        for word in words:
            test_line = ' '.join(current_line + [word]) if current_line else word
            # Calculate width of the test line
            bbox = font.getbbox(test_line)
            width = bbox[2] - bbox[0]
            if width <= max_width:
                current_line.append(word)
            else:
                lines.append(' '.join(current_line))
                current_line = [word]
        if current_line:
            lines.append(' '.join(current_line))
    return lines

def main():
    parser = argparse.ArgumentParser(description="Render verbatim prompt text into a high-contrast monochrome OCR-friendly PNG.")
    parser.add_argument("-i", "--input", required=True, help="Path to input text file containing the prompt")
    parser.add_argument("-o", "--output", required=True, help="Path to save output PNG image")
    parser.add_argument("-w", "--width", type=int, default=800, help="Width of the output image in pixels")
    parser.add_argument("-s", "--font-size", type=int, default=16, help="Font size to use")
    
    args = parser.parse_args()

    if not os.path.exists(args.input):
        print(f"Error: Input file {args.input} does not exist.", file=sys.stderr)
        sys.exit(1)

    try:
        with open(args.input, "r", encoding="utf-8") as f:
            text = f.read()
    except Exception as e:
        print(f"Error reading input file: {e}", file=sys.stderr)
        sys.exit(1)

    try:
        from PIL import Image, ImageDraw, ImageFont
    except ImportError:
        print("Error: Pillow library is required. Install with: pip install Pillow", file=sys.stderr)
        sys.exit(1)

    # 1. Load font (prefer a standard macOS monospace/sans font for OCR clarity)
    font = None
    font_paths = [
        "/System/Library/Fonts/Supplemental/Courier New.ttf",
        "/System/Library/Fonts/Monaco.ttf",
        "/System/Library/Fonts/Supplemental/Arial.ttf",
        "/Library/Fonts/Arial.ttf"
    ]
    for fp in font_paths:
        if os.path.exists(fp):
            try:
                font = ImageFont.truetype(fp, args.font_size)
                break
            except Exception:
                pass
    
    if font is None:
        print("Warning: Standard fonts not found. Falling back to default font.", file=sys.stderr)
        font = ImageFont.load_default()

    # 2. Calculate text layout & wrapping
    padding = 20
    draw_width = args.width - (padding * 2)
    
    wrapped_lines = wrap_text(text, font, draw_width)
    
    # Calculate line height
    test_bbox = font.getbbox("Ayj|")
    line_height = int((test_bbox[3] - test_bbox[1]) * 1.3) # 1.3x line spacing
    if line_height <= 0:
        line_height = args.font_size + 4

    total_height = padding * 2 + len(wrapped_lines) * line_height

    # 3. Create 1-bit image (mode "1")
    # 255 = White background, 0 = Black text
    img = Image.new("1", (args.width, total_height), 255)
    draw = ImageDraw.Draw(img)

    # 4. Draw text
    y = padding
    for line in wrapped_lines:
        if line:
            draw.text((padding, y), line, fill=0, font=font)
        y += line_height

    # 5. Save optimized image
    try:
        out_dir = os.path.dirname(args.output)
        if out_dir:
            os.makedirs(out_dir, exist_ok=True)
            
        img.save(args.output, format="PNG", optimize=True)
        print(f"Success: Rendered prompt text into {args.output}")
        print(f"  Lines: {len(wrapped_lines)} | Dimensions: {args.width}x{total_height}")
        print(f"  File Size: {os.path.getsize(args.output) / 1024:.2f} KB")
    except Exception as e:
        print(f"Error saving image: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()