src/compliment_forest/style_data.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

from .backends.image import STYLE_PROFILES, compose_flux_prompt

SceneCategory = Literal["animal", "human", "object", "environment"]


@dataclass(frozen=True)
class ForestScene:
    slug: str
    category: SceneCategory
    prompt: str


FOREST_SCENES = (
    ForestScene(
        "fox-threshold",
        "animal",
        "a gentle red fox pausing at the edge of a fern-lined path",
    ),
    ForestScene(
        "listening-owl",
        "animal",
        "a round barn owl listening from a low mossy branch",
    ),
    ForestScene(
        "steady-deer",
        "animal",
        "a young deer standing calmly between silver birch trees",
    ),
    ForestScene(
        "brave-snail",
        "animal",
        "a tiny snail crossing a dew-covered fern at dawn",
    ),
    ForestScene(
        "singing-wren",
        "animal",
        "a small wren singing beside loose woodland flowers",
    ),
    ForestScene(
        "river-otter",
        "animal",
        "a river otter holding one smooth stone beside quiet reeds",
    ),
    ForestScene(
        "thoughtful-badger",
        "animal",
        "a thoughtful badger beside a lantern-shaped mushroom",
    ),
    ForestScene(
        "patient-hare",
        "animal",
        "a patient brown hare resting beneath arching grasses",
    ),
    ForestScene(
        "moonlit-moth",
        "animal",
        "a luna moth hovering near moonlit foxgloves",
    ),
    ForestScene(
        "walking-turtle",
        "animal",
        "a small woodland turtle moving between clover and stones",
    ),
    ForestScene(
        "person-open-window",
        "human",
        "an adult seen from behind opening a window to pale morning light",
    ),
    ForestScene(
        "person-blank-notebook",
        "human",
        "an adult seated at a wooden desk with an open blank notebook",
    ),
    ForestScene(
        "person-forked-path",
        "human",
        "a small human figure viewed from behind at a gentle fork in a path",
    ),
    ForestScene(
        "person-train-platform",
        "human",
        "a quiet adult figure waiting on a misty train platform with one bag",
    ),
    ForestScene(
        "person-moving-box",
        "human",
        "an adult carrying one moving box toward a sunlit doorway",
    ),
    ForestScene(
        "person-footbridge",
        "human",
        "a side-view figure taking one step across a narrow wooden footbridge",
    ),
    ForestScene(
        "person-doorway",
        "human",
        "a calm adult silhouette standing in an open doorway between two rooms",
    ),
    ForestScene(
        "person-seedling",
        "human",
        "hands gently watering a small seedling on a windowsill",
    ),
    ForestScene(
        "person-rain-shelter",
        "human",
        "an adult seen from the side resting on a bench beneath a rain shelter",
    ),
    ForestScene(
        "person-dawn-hill",
        "human",
        "a distant human figure standing on a low hillside at dawn",
    ),
    ForestScene(
        "lantern-crossroads",
        "object",
        "a small glowing lantern placed where two woodland paths meet",
    ),
    ForestScene(
        "map-compass",
        "object",
        "an unfolded map and simple compass resting on a wooden table",
    ),
    ForestScene(
        "open-notebook",
        "object",
        "an open blank notebook beside a pencil and one pressed leaf",
    ),
    ForestScene(
        "stepping-stones",
        "object",
        "four smooth stepping stones crossing a narrow stream",
    ),
    ForestScene(
        "warm-cup",
        "object",
        "a warm ceramic cup sending a thin curl of steam into morning light",
    ),
    ForestScene(
        "woven-thread",
        "object",
        "loose green and gold threads gradually woven into one calm pattern",
    ),
    ForestScene(
        "key-and-door",
        "object",
        "a simple brass key resting beside a small unopened wooden door",
    ),
    ForestScene(
        "paper-boat",
        "object",
        "a single paper boat floating on still water beneath willow reflections",
    ),
    ForestScene(
        "balanced-stones",
        "object",
        "three imperfect river stones balanced beside soft grasses",
    ),
    ForestScene(
        "empty-chair-light",
        "object",
        "an empty wooden chair in a quiet patch of warm window light",
    ),
    ForestScene(
        "winding-path",
        "environment",
        "a winding path disappearing gently through tall ferns and morning mist",
    ),
    ForestScene(
        "river-crossing",
        "environment",
        "a shallow river crossing with stones visible beneath clear water",
    ),
    ForestScene(
        "room-at-dawn",
        "environment",
        "a quiet room at dawn with curtains moving beside an open window",
    ),
    ForestScene(
        "city-garden",
        "environment",
        "a small green garden between quiet city buildings after rain",
    ),
    ForestScene(
        "misty-platform",
        "environment",
        "an empty train platform fading softly into early morning mist",
    ),
    ForestScene(
        "clearing-after-rain",
        "environment",
        "a forest clearing just after rain with one bright opening in the clouds",
    ),
    ForestScene(
        "hillside-trail",
        "environment",
        "a gradual hillside trail curving toward a pale open horizon",
    ),
    ForestScene(
        "staircase-light",
        "environment",
        "a simple staircase with warm light falling across the next three steps",
    ),
    ForestScene(
        "canopy-opening",
        "environment",
        "a dark green canopy opening into a circle of soft sky",
    ),
    ForestScene(
        "shoreline-horizon",
        "environment",
        "a calm shoreline where fading clouds meet a wide quiet horizon",
    ),
)

# Compatibility alias for callers that used the v1 name.
FOREST_SUBJECTS = FOREST_SCENES

TRAINED_STYLE_IDS = (
    "watercolor",
    "paper_cut",
    "moonlit_gouache",
    "botanical_ink",
)


def build_style_records(
    *,
    samples_per_style: int = 40,
    base_seed: int = 9000,
) -> list[dict[str, str | int]]:
    if not 1 <= samples_per_style <= len(FOREST_SCENES):
        raise ValueError(f"samples_per_style must be between 1 and {len(FOREST_SCENES)}")

    records: list[dict[str, str | int]] = []
    for style_offset, style in enumerate(TRAINED_STYLE_IDS):
        profile = STYLE_PROFILES[style]
        for scene_index, scene in enumerate(FOREST_SCENES[:samples_per_style]):
            seed = base_seed + style_offset * 1000 + scene_index
            prompt = compose_flux_prompt(
                scene.prompt,
                style=style,  # type: ignore[arg-type]
                seed=seed,
            )
            records.append(
                {
                    "style": style,
                    "trigger": profile.trigger,
                    "category": scene.category,
                    "subject": scene.slug,
                    "seed": seed,
                    "prompt": prompt,
                    "text": (
                        f"{profile.trigger}, {scene.prompt}, "
                        f"{profile.label.lower()} storybook scene"
                    ),
                    "file_name": (
                        f"{scene_index:03d}-{scene.category}-{scene.slug}.png"
                    ),
                }
            )
    return records