src/image_utils.py

"""Greenscreen chroma key for #00FF00 bait/fish sprites."""

from __future__ import annotations

from io import BytesIO
from typing import Any

import numpy as np
from PIL import Image


def chroma_key_green(
    image: Image.Image,
    *,
    key: tuple[int, int, int] = (0, 255, 0),
    tolerance: int = 90,
    feather: int = 40,
    spill_suppression: float = 0.5,
) -> Image.Image:
    """Replace green-screen pixels with transparency (vectorized)."""
    img = image.convert("RGBA")
    arr = np.asarray(img).astype(np.float32)
    r, g, b = arr[..., 0], arr[..., 1], arr[..., 2]

    kr, kg, kb = key
    dist = np.sqrt((r - kr) ** 2 + (g - kg) ** 2 + (b - kb) ** 2)

    greenness = g - np.maximum(r, b)
    green_dominant = (greenness > 60) & (g > 120)

    alpha = np.clip((dist - tolerance) / max(feather, 1), 0.0, 1.0) * 255.0
    alpha = np.where(green_dominant & (dist < tolerance + feather * 2), 0.0, alpha)
    arr[..., 3] = np.minimum(arr[..., 3], alpha)

    if spill_suppression > 0:
        spill = np.clip(greenness, 0, None) * spill_suppression
        arr[..., 1] = np.clip(g - spill, 0, 255)

    return Image.fromarray(arr.astype(np.uint8), "RGBA")


def trim_transparent(image: Image.Image, padding: int = 4) -> Image.Image:
    """Crop away fully transparent borders so sprites sit tight in the UI."""
    img = image.convert("RGBA")
    alpha = np.asarray(img)[..., 3]
    rows = np.any(alpha > 8, axis=1)
    cols = np.any(alpha > 8, axis=0)
    if not rows.any() or not cols.any():
        return img
    top, bottom = np.argmax(rows), len(rows) - np.argmax(rows[::-1])
    left, right = np.argmax(cols), len(cols) - np.argmax(cols[::-1])
    top = max(0, top - padding)
    left = max(0, left - padding)
    bottom = min(img.height, bottom + padding)
    right = min(img.width, right + padding)
    return img.crop((left, top, right, bottom))


def image_to_png_bytes(image: Image.Image) -> bytes:
    buf = BytesIO()
    image.save(buf, format="PNG")
    return buf.getvalue()


def bait_prompt(bait: str) -> str:
    return (
        f'A "{bait}" cartoon game item sprite, Stardew Valley chibi pixel art style, '
        "single object centered, bold flat colours, solid bright green #00FF00 "
        "greenscreen background, flat uniform background colour, no shadows on background"
    )


def _fish_label(fish: dict[str, Any], bait: str = "") -> str:
    name = fish.get("name", "fish")
    appearance = fish.get(
        "appearance",
        f"cute cartoon fish themed after {bait}" if bait else "cute cartoon fish",
    )
    if bait:
        return f"{name} (bait was {bait}): {appearance}"
    return f"{name}: {appearance}"


def fish_underwater_prompt(fish: dict[str, Any], bait: str = "") -> str:
    label = _fish_label(fish, bait)
    return (
        f"Stardew Valley cartoon pixel art sprite, side view of one silly fish — {label}. "
        "Chibi proportions, big eye, bold flat colours, exaggerated cute shape. "
        "The fish body clearly mimics the bait theme. Simple underwater bubbles, "
        "soft colours. NOT photorealistic, NOT realistic."
    )


def fish_edit_prompt(fish: dict[str, Any], bait: str = "") -> str:
    label = _fish_label(fish, bait)
    return (
        f"Keep this exact cartoon fish sprite ({label}) unchanged. Replace background "
        "with solid bright green #00FF00 greenscreen. No water. Same pixel art style."
    )


def fish_greenscreen_prompt(fish: dict[str, Any], bait: str = "") -> str:
    label = _fish_label(fish, bait)
    return (
        f"Stardew Valley cartoon pixel art sprite, side view — {label}. "
        "Chibi fish, bait-themed body, big eye, fins and tail. "
        "Solid bright green #00FF00 greenscreen background. NOT photorealistic."
    )