src/tool.py

"""AI art generator: prompt assembly, generation, file handling."""

import html
import os
import random
import tempfile
from pathlib import Path

from .shared import image_client
from .shared.ai_client import AIUnavailable
from .shared.ratelimit import RateLimiter

# Images cost real money per generation (~$0.05-0.09 on krea-2-large), so this
# Space runs tighter limits than the text tools. Overridable via Space vars.
limiter = RateLimiter(
    per_ip_limit=int(os.environ.get("LF_IP_LIMIT", "5")),
    window_seconds=int(os.environ.get("LF_IP_WINDOW_SECONDS", "3600")),
    daily_cap=int(os.environ.get("LF_DAILY_CAP", "200")),
)

RACES = [
    "Human", "Elf", "Half-Elf", "Dwarf", "Halfling", "Dragonborn", "Tiefling",
    "Half-Orc", "Orc", "Gnome", "Goliath", "Aasimar", "Tabaxi", "Kenku", "Firbolg",
]
CLASSES = [
    "Fighter", "Wizard", "Rogue", "Cleric", "Ranger", "Paladin", "Barbarian",
    "Bard", "Druid", "Monk", "Sorcerer", "Warlock", "Artificer",
]
STYLES = {
    "Classic fantasy oil painting":
        "classic high-fantasy oil painting, rich color, dramatic lighting, detailed brushwork",
    "Character-sheet portrait":
        "clean character-sheet portrait illustration, neutral parchment background, crisp linework, soft studio lighting",
    "Dark gritty fantasy":
        "dark gritty fantasy illustration, moody chiaroscuro lighting, weathered textures, muted palette",
    "Painterly watercolor":
        "loose painterly watercolor illustration, soft washes, ink linework, white paper background",
    "Comic / animated":
        "stylized comic-book illustration, bold outlines, cel shading, expressive pose",
    "Isometric token art":
        "top-down VTT token art, centered figure, circular composition, high contrast against plain background",
    "Pencil sketch":
        "detailed graphite pencil sketch, cross-hatching, sketchbook style, monochrome",
}
FRAMINGS = {
    "Portrait (head & shoulders)": "head-and-shoulders portrait, face in sharp focus",
    "Half body": "half-body composition from the waist up, hands visible",
    "Full body": "full-body character illustration, complete costume and gear visible",
    "Round token": "circular VTT token crop, figure centered and facing forward",
}

MAX_DESC = 600
MAX_FIELD = 120

PROMPT_SUFFIX = (
    "tasteful fantasy character artwork, high detail, no text, no watermark, "
    "no signature, no frame"
)


def randomize():
    return (
        random.choice(RACES),
        random.choice(CLASSES),
        random.choice(list(STYLES)),
        random.choice(list(FRAMINGS)),
    )


def _clip(value: str, limit: int) -> str:
    return (value or "").strip()[:limit]


def build_prompt(description, race, char_class, style, framing, palette) -> str:
    description = _clip(description, MAX_DESC)
    race = _clip(race, MAX_FIELD)
    char_class = _clip(char_class, MAX_FIELD)
    palette = _clip(palette, MAX_FIELD)
    style_text = STYLES.get(style, STYLES["Classic fantasy oil painting"])
    framing_text = FRAMINGS.get(framing, FRAMINGS["Portrait (head & shoulders)"])

    subject_bits = [b for b in (race, char_class) if b]
    subject = " ".join(subject_bits) if subject_bits else "adventurer"
    if description:
        subject = f"{subject}: {description}"
    else:
        subject = f"a heroic {subject} with a distinctive, memorable look"

    parts = [
        f"Dungeons and Dragons character art of a {subject}",
        framing_text,
        style_text,
    ]
    if palette:
        parts.append(f"color palette: {palette}")
    parts.append(PROMPT_SUFFIX)
    return ". ".join(parts)


def generate(description, race, char_class, style, framing, palette, request=None):
    """Returns (image_path, download_path, error_message)."""
    allowed, message = limiter.check(request)
    if not allowed:
        return None, None, message

    prompt = build_prompt(description, race, char_class, style, framing, palette)
    try:
        image_bytes, media_type = image_client.generate_image(prompt)
    except AIUnavailable as exc:
        return None, None, str(exc)

    ext = {"image/png": "png", "image/jpeg": "jpg", "image/webp": "webp"}.get(media_type, "png")
    path = Path(tempfile.mkdtemp()) / f"dnd-character-art.{ext}"
    path.write_bytes(image_bytes)
    return str(path), str(path), None


def render_error(message: str) -> str:
    return (
        '<div class="lf-output" role="alert" style="border-color:#A7343B;">'
        f"<p>{html.escape(message)}</p></div>"
    )