#!/usr/bin/env python3
"""Generate placeholder preview images for ComfyUI workflow templates."""

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

PREVIEW_WIDTH = 512
PREVIEW_HEIGHT = 288  # 16:9 aspect ratio

# Category colors based on filename patterns
# Order matters: more specific patterns first, generic ones last
CATEGORY_COLORS = {
    # NSFW patterns (check first - most specific)
    "lustify": "#9B2C2C",      # Dark red - NSFW
    "realvis": "#9B2C2C",      # Dark red - NSFW
    "pony": "#9B2C2C",         # Dark red - NSFW
    "nsfw": "#9B2C2C",         # Dark red - NSFW
    # Advanced patterns
    "animatediff": "#4ad9d9",  # Cyan - advanced
    "controlnet": "#4ad9d9",   # Cyan - advanced
    "batch": "#4ad9d9",        # Cyan - advanced
    # Upscaling
    "upscale": "#9b4ad9",      # Purple - upscaling
    # Generic type patterns (check last)
    "t2i": "#4a90d9",          # Blue - text-to-image
    "i2v": "#d94a4a",          # Red - image-to-video
    "i2i": "#4ad96b",          # Green - image-to-image
    "t2m": "#d9a04a",          # Orange - text-to-music
    "m2m": "#d9a04a",          # Orange - music-to-music
}
DEFAULT_COLOR = "#2d2d2d"


def get_color_for_workflow(filename: str) -> str:
    """Determine background color based on filename."""
    lower = filename.lower()
    for pattern, color in CATEGORY_COLORS.items():
        if pattern in lower:
            return color
    return DEFAULT_COLOR


def humanize_filename(filename: str) -> str:
    """Convert filename to readable title."""
    # Remove extension and common suffixes
    name = filename.replace('.json', '')
    name = name.replace('-production-v1', '')
    name = name.replace('-production', '')
    name = name.replace('-v1', '')
    # Replace separators with spaces
    name = name.replace('-', ' ').replace('_', ' ')
    # Title case
    return name.title()


def generate_preview(json_path: Path, output_path: Path):
    """Generate a placeholder preview image with workflow title."""
    filename = json_path.name
    bg_color = get_color_for_workflow(filename)
    title = humanize_filename(filename)

    # Create image
    img = Image.new('RGB', (PREVIEW_WIDTH, PREVIEW_HEIGHT), bg_color)
    draw = ImageDraw.Draw(img)

    # Try to load fonts
    try:
        title_font = ImageFont.truetype(
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 24
        )
    except OSError:
        title_font = ImageFont.load_default()

    # Word wrap title
    words = title.split()
    lines = []
    current_line = ""
    for word in words:
        test = f"{current_line} {word}".strip()
        if len(test) > 28:
            if current_line:
                lines.append(current_line)
            current_line = word
        else:
            current_line = test
    if current_line:
        lines.append(current_line)

    # Draw centered text
    y_offset = PREVIEW_HEIGHT // 2 - (len(lines) * 30) // 2
    for line in lines:
        bbox = draw.textbbox((0, 0), line, font=title_font)
        text_width = bbox[2] - bbox[0]
        x = (PREVIEW_WIDTH - text_width) // 2
        draw.text((x, y_offset), line, fill="white", font=title_font)
        y_offset += 30

    img.save(output_path, 'JPEG', quality=85)


def main():
    workflows_dir = Path(__file__).parent.parent / 'example_workflows'

    count = 0
    for json_file in sorted(workflows_dir.glob('*.json')):
        preview_path = json_file.with_suffix('.jpg')

        if preview_path.exists():
            print(f"Skipping (exists): {json_file.name}")
            continue

        generate_preview(json_file, preview_path)
        print(f"Generated: {preview_path.name}")
        count += 1

    print(f"\nDone! Generated {count} preview images.")


if __name__ == '__main__':
    main()