image.py

#!/usr/bin/env python3
"""
Usage:
  python image.py <folder_key_or_group> [allowed_degrees]

Examples:
  python image.py core_conflict
  python image.py conflict 0,180
  python image.py anchor 0,90,180,270

Behavior:
- If <folder_key_or_group> is an exact key in FOLDERS, use that folder.
- If it is one of: agent, anchor, aspect, conflict, engine
  pick a random folder whose key contains that token.

Image selection:
- Only select source images that do NOT already have a rotation suffix in the filename.

Rotation / output:
- Choose a random rotation from the allowed degree list (defaults to 0,90,180,270).
- If rotation is 0, return the original file (no copy).
- If rotation is non-zero:
  - If the rotated filename already exists, return it (no overwrite).
  - Else, create it and return it.

Output:
Print a Markdown image link using a file:// URI.
"""

from __future__ import annotations

import argparse
import random
import sys
from pathlib import Path
from typing import List, Tuple
from urllib.parse import quote

from PIL import Image

# Map folder keys to folder paths (edit these)
FOLDERS: Dict[str, str] = {
    "core_agent": r"C:\story_engine\0-Core\Core Agent 36",
    "core_anchor": r"C:\story_engine\0-Core\Core Anchor 36",
    "core_aspect": r"C:\story_engine\0-Core\Core Aspect 36",
    "core_conflict": r"C:\story_engine\0-Core\Core Conflict 36",
    "core_engine": r"C:\story_engine\0-Core\Core Engine 36",
    "fantasy_agent": r"C:\story_engine\1-Fantasy\Fantasy Agent 12",
    "fantasy_anchor": r"C:\story_engine\1-Fantasy\Fantasy Anchor 12",
    "fantasy_aspect": r"C:\story_engine\1-Fantasy\Fantasy Aspect 12",
    "fantasy_conflict": r"C:\story_engine\1-Fantasy\Fantasy Conflict 12",
    "fantasy_engine": r"C:\story_engine\1-Fantasy\Fantasy Engine 12",
    "horror_agent": r"C:\story_engine\2-Horror\Horror Agent 12",
    "horror_anchor": r"C:\story_engine\2-Horror\Horror Anchor 12",
    "horror_aspect": r"C:\story_engine\2-Horror\Horror Aspect 12",
    "horror_conflict": r"C:\story_engine\2-Horror\Horror Conflict 12",
    "horror_engine": r"C:\story_engine\2-Horror\Horror Engine 12",
    "sci_fi_agent": r"C:\story_engine\3-Sci-Fi\Sci-Fi Agent 12",
    "sci_fi_anchor": r"C:\story_engine\3-Sci-Fi\Sci-Fi Anchor 12",
    "sci_fi_aspect": r"C:\story_engine\3-Sci-Fi\Sci-Fi Aspect 12",
    "sci_fi_conflict": r"C:\story_engine\3-Sci-Fi\Sci-Fi Conflict 12",
    "sci_fi_engine": r"C:\story_engine\3-Sci-Fi\Sci-Fi Engine 12",
    "mystery_agent": r"C:\story_engine\4-Mystery\Mystery Agent 12",
    "mystery_anchor": r"C:\story_engine\4-Mystery\Mystery Anchor 12",
    "mystery_aspect": r"C:\story_engine\4-Mystery\Mystery Aspect 12",
    "mystery_conflict": r"C:\story_engine\4-Mystery\Mystery Conflict 12",
    "mystery_engine": r"C:\story_engine\4-Mystery\Mystery Engine 12",
    "backstories_conflict": r"C:\story_engine\5-Backstories\Backstories Conflict",
    "backstories_engine": r"C:\story_engine\5-Backstories\Backstories Engine",
    "items_anchor": r"C:\story_engine\6-Items\Items Anchor 12",
    "items_aspect": r"C:\story_engine\6-Items\Items Aspect 48",
    "founder_agent": r"C:\story_engine\7-Founder\Founder Agent 30",
    "founder_anchor": r"C:\story_engine\7-Founder\Founder Anchor 30",
    "founder_aspect": r"C:\story_engine\7-Founder\Founder Aspect 24",
    "founder_conflict": r"C:\story_engine\7-Founder\Founder Conflict 12",
    "founder_engine": r"C:\story_engine\7-Founder\Founder Engine 12",
    "dreamer_agent": r"C:\story_engine\8-Dreamer\Dreamer Agent 24",
    "dreamer_anchor": r"C:\story_engine\8-Dreamer\Dreamer Anchor 24",
    "dreamer_aspect": r"C:\story_engine\8-Dreamer\Dreamer Aspect 24",
    "dreamer_conflict": r"C:\story_engine\8-Dreamer\Dreamer Conflict 24",
    "dreamer_engine": r"C:\story_engine\8-Dreamer\Dreamer Engine 24",
}

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"}
GROUP_TOKENS = {"agent", "anchor", "aspect", "conflict", "engine"}

# These are the only suffixes this script considers "rotation suffixes"
ROT_SUFFIXES = {
    "rotate_left",
    "rotate_right",
    "upside_down",
}


def parse_allowed_degrees(s: str | None) -> List[int]:
    if not s:
        return [0, 90, 180, 270]
    parts = [p.strip() for p in s.split(",") if p.strip()]
    if not parts:
        return [0, 90, 180, 270]

    allowed = []
    for p in parts:
        try:
            deg = int(p)
        except ValueError as e:
            raise ValueError(f"Invalid degree '{p}'. Expected integers like 0,90,180,270.") from e
        if deg % 90 != 0:
            raise ValueError(f"Invalid degree '{deg}'. Must be a multiple of 90.")
        deg = deg % 360
        allowed.append(deg)

    # de-dup but preserve order
    out: List[int] = []
    seen = set()
    for d in allowed:
        if d not in seen:
            out.append(d)
            seen.add(d)
    return out


def degree_to_suffix(deg: int) -> str:
    deg = deg % 360
    if deg == 0:
        return ""
    if deg == 90:
        return "rotate_right"
    if deg == 180:
        return "upside_down"
    if deg == 270:
        return "rotate_left"
    raise ValueError(f"Unsupported rotation: {deg}")


def is_rotated_variant(path: Path) -> bool:
    """
    Returns True if filename ends with _<rotation_suffix> before extension.
    Examples:
      Core-Conflict (8)_upside_down.png -> True
      Core-Conflict (8)_rotate_left__2.png -> True
      Core-Conflict (8).png -> False
    """
    stem = path.stem  # no extension
    for suf in ROT_SUFFIXES:
        if stem.endswith(f"_{suf}"):
            return True
        if f"_{suf}__" in stem and stem.split(f"_{suf}__", 1)[0] != stem:
            # matches patterns like ..._<suf>__2
            return True
    return False


def list_source_images(folder: Path) -> List[Path]:
    if not folder.exists() or not folder.is_dir():
        raise FileNotFoundError(f"Folder does not exist or is not a directory: {folder}")
    imgs = []
    for p in folder.iterdir():
        if not p.is_file():
            continue
        if p.suffix.lower() not in IMAGE_EXTS:
            continue
        if is_rotated_variant(p):
            continue
        imgs.append(p)
    return sorted(imgs)


def rotate_image(src: Path, degrees_cw: int) -> Image.Image:
    img = Image.open(src)
    img.load()

    deg = degrees_cw % 360
    if deg == 0:
        return img
    if deg == 90:
        return img.transpose(Image.Transpose.ROTATE_270)
    if deg == 180:
        return img.transpose(Image.Transpose.ROTATE_180)
    if deg == 270:
        return img.transpose(Image.Transpose.ROTATE_90)

    raise ValueError(f"Unsupported rotation: {degrees_cw}")


def rotated_out_path(src: Path, suffix: str) -> Path:
    return src.with_name(f"{src.stem}_{suffix}{src.suffix}")


def save_rotated_if_needed(src: Path, degrees_cw: int) -> Path:
    deg = degrees_cw % 360
    if deg == 0:
        return src

    suffix = degree_to_suffix(deg)
    out_path = rotated_out_path(src, suffix)

    # If it already exists, do not overwrite, just return it
    if out_path.exists():
        return out_path

    with Image.open(src) as img:
        img.load()
        rotated = rotate_image(src, deg)

        fmt = (img.format or "").upper()
        save_kwargs = {}
        if fmt == "JPEG":
            save_kwargs["quality"] = 95
            save_kwargs["subsampling"] = 0
            save_kwargs["optimize"] = True

        if fmt in {"JPEG", "PNG", "WEBP", "BMP", "TIFF"}:
            rotated.save(out_path, format=fmt, **save_kwargs)
        else:
            rotated.save(out_path, format="PNG")

    return out_path


def path_to_file_uri(path: Path) -> str:
    resolved = path.resolve()
    posix_path = resolved.as_posix()
    encoded_path = quote(posix_path, safe="/:")
    if resolved.drive:
        return f"file:///{encoded_path}"
    return f"file://{encoded_path}"


def resolve_folder_key(user_arg: str) -> str:
    arg_l = user_arg.strip().lower()

    if user_arg in FOLDERS:
        return user_arg

    if arg_l in GROUP_TOKENS:
        candidates = [k for k in FOLDERS.keys() if arg_l in k.lower()]
        if not candidates:
            raise KeyError(f"No folder keys contain token '{arg_l}'.")
        return random.choice(candidates)

    raise KeyError(f"Unknown folder key/group: '{user_arg}'.")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("folder_key", help="Exact key in FOLDERS, or one of: agent/anchor/aspect/conflict/engine")
    parser.add_argument(
        "allowed_degrees",
        nargs="?",
        default=None,
        help="Comma-separated degrees: 0,90,180,270 (default: 0,90,180,270)",
    )
    parser.add_argument("--seed", type=int, default=None)
    args = parser.parse_args()

    if args.seed is not None:
        random.seed(args.seed)

    try:
        allowed_degrees = parse_allowed_degrees(args.allowed_degrees)
    except ValueError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        return 2

    try:
        chosen_key = resolve_folder_key(args.folder_key)
    except KeyError as e:
        keys = ", ".join(sorted(FOLDERS.keys()))
        groups = ", ".join(sorted(GROUP_TOKENS))
        print(f"ERROR: {e}", file=sys.stderr)
        print(f"Available groups: {groups}", file=sys.stderr)
        print(f"Available exact keys: {keys}", file=sys.stderr)
        return 2

    folder = Path(FOLDERS[chosen_key]).expanduser().resolve()

    try:
        images = list_source_images(folder)
    except FileNotFoundError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        return 2

    if not images:
        print(f"ERROR: No non-rotated source images found in: {folder}", file=sys.stderr)
        return 2

    src = random.choice(images)
    deg = random.choice(allowed_degrees)
    out = save_rotated_if_needed(src, deg)

    uri = path_to_file_uri(out)
    print(f"![]({uri})")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())