import argparse
import re
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont

# --- Configuration ---
PADDING = 40
FONT_SIZE = 48
FONT_COLOR = "black"
ARROW_COLOR = "black"
BACKGROUND_COLOR = "white"
ARROW_WIDTH_RATIO = 0.3
ARROW_HEIGHT_RATIO = 0.1

def parse_filename(filepath: Path):
    """Extracts the step number and name from a filename."""
    # Pattern for standard steps like '..._02_log_exposure_RGB.jpg'
    match = re.search(r'_(\d+)_([a-zA-Z0-9_]+?)_RGB\.', filepath.name)
    if match:
        step_name = match.group(2).replace('_', ' ').title()
        return int(match.group(1)), step_name

    # Fallback pattern for the first input image like '..._input_linear_sRGB.jpg'
    match_input = re.search(r'_input_([a-zA-Z0-9_]+?)_sRGB\.', filepath.name)
    if match_input:
        step_name = f"Input {match_input.group(1).replace('_', ' ').title()}"
        return 0, step_name
        
    # If no pattern matches, return a generic name
    return 999, filepath.stem.replace('_', ' ').title()

def create_arrow(width: int, height: int, color: str) -> Image.Image:
    """Creates a right-pointing arrow image with a transparent background."""
    arrow_img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    draw = ImageDraw.Draw(arrow_img)
    
    shaft_width = width * 0.7
    rect_start_y = (height // 2) - (height // 10)
    rect_height = max(1, height // 5)
    
    draw.rectangle([(0, rect_start_y), (shaft_width, rect_start_y + rect_height)], fill=color)
    draw.polygon([(shaft_width, 0), (width, height // 2), (shaft_width, height)], fill=color)
    
    return arrow_img

def main():
    parser = argparse.ArgumentParser(
        description="Create a visual pipeline of image processing steps with diffs.",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument("input_dir", type=str, help="Directory containing the input and diff images.")
    parser.add_argument("output_file", type=str, help="Path for the final combined image.")
    parser.add_argument(
        "--scale", type=float, default=0.25,
        help="Scale factor for the main pipeline images (e.g., 0.25 for 25%%). Default is 0.25."
    )
    parser.add_argument(
        "--diff-scale", type=float, default=0.15,
        help="Scale factor for the smaller diff images. Default is 0.15."
    )
    args = parser.parse_args()

    input_path = Path(args.input_dir)
    if not input_path.is_dir():
        print(f"Error: Input directory '{args.input_dir}' not found.")
        return

    # 1. Find and sort all images
    # --- THIS IS THE FIX ---
    # Use a more general glob to find all .jpg files for the pipeline
    pipeline_images_raw = list(input_path.glob("*.jpg"))
    diff_images_raw = list(input_path.glob("diff_*_RGB.png"))

    if not pipeline_images_raw:
        print("Error: No pipeline images (*.jpg) found in the directory.")
        return

    # Sort files alphabetically by their full name, which mimics 'ls -1' behavior.
    pipeline_images_sorted = sorted(pipeline_images_raw)
    diff_images_sorted = sorted(diff_images_raw)
    
    print("Found and sorted the following pipeline images:")
    for p in pipeline_images_sorted:
        print(f" - {p.name}")

    # 2. Prepare images and assets
    with Image.open(pipeline_images_sorted[0]) as img:
        orig_w, orig_h = img.size
    
    img_w, img_h = int(orig_w * args.scale), int(orig_h * args.scale)
    diff_w, diff_h = int(orig_w * args.diff_scale), int(orig_h * args.scale)

    pipeline_images = [Image.open(p).resize((img_w, img_h), Image.Resampling.LANCZOS) for p in pipeline_images_sorted]
    diff_images = [Image.open(p).resize((diff_w, diff_h), Image.Resampling.LANCZOS) for p in diff_images_sorted]

    arrow_w, arrow_h = int(img_w * ARROW_WIDTH_RATIO), int(img_h * ARROW_HEIGHT_RATIO)
    arrow = create_arrow(arrow_w, arrow_h, ARROW_COLOR)
    
    try:
        font = ImageFont.truetype("arial.ttf", FONT_SIZE)
    except IOError:
        print("Arial font not found, using default font.")
        font = ImageFont.load_default()

    # 3. Calculate canvas size
    num_steps = len(pipeline_images)
    gap_width = max(arrow_w, diff_w) + PADDING
    total_width = (num_steps * img_w) + ((num_steps - 1) * gap_width) + (2 * PADDING)
    total_height = PADDING + FONT_SIZE + PADDING + img_h + PADDING + diff_h + PADDING

    # 4. Create the final canvas and draw everything
    canvas = Image.new("RGB", (total_width, total_height), BACKGROUND_COLOR)
    draw = ImageDraw.Draw(canvas)

    y_text = PADDING
    y_pipeline = y_text + FONT_SIZE + PADDING
    y_arrow = y_pipeline + (img_h // 2) - (arrow_h // 2)
    y_diff = y_pipeline + img_h + PADDING
    current_x = PADDING

    for i, p_img in enumerate(pipeline_images):
        canvas.paste(p_img, (current_x, y_pipeline))
        
        _, step_name = parse_filename(pipeline_images_sorted[i])
        if step_name:
            text_bbox = draw.textbbox((0, 0), step_name, font=font)
            text_w = text_bbox[2] - text_bbox[0]
            draw.text((current_x + (img_w - text_w) // 2, y_text), step_name, font=font, fill=FONT_COLOR)

        if i < num_steps - 1:
            gap_start_x = current_x + img_w
            
            arrow_x = gap_start_x + (gap_width - arrow_w) // 2
            canvas.paste(arrow, (arrow_x, y_arrow), mask=arrow)
            
            if i < len(diff_images):
                diff_x = gap_start_x + (gap_width - diff_w) // 2
                canvas.paste(diff_images[i], (diff_x, y_diff))
            
            current_x += img_w + gap_width

    # 5. Save the final image
    canvas.save(args.output_file, quality=95)
    print(f"\nPipeline image successfully created at '{args.output_file}'")

if __name__ == "__main__":
    main()