"""ShotCraft Stage 2 — Frame Generator (FLUX.1-schnell, 12B). Owner: Rafal.

Rendering runs on the Modal backend (model_runtime.flux_generate_batch):
the whole reel is one backend call; a single-frame regen is a 1-element batch.
"""
from __future__ import annotations

from model_runtime import flux_generate_batch
from schemas import ConceptPackage, STYLE_SUFFIXES

MODEL_ID = "black-forest-labs/FLUX.1-schnell"
NUM_STEPS = 4            # schnell sweet spot, fixed server-side (NFR-1)
SEED_BASE = 42           # FR-2.3: deterministic seed base
SIZES = {"1:1": (1024, 1024), "16:9": (1024, 576)}  # FR-2.6

def build_prompt(shot, style_preset: str, product_desc: str = "") -> str:
    """FR-2.3: prompt + palette keywords + shared style suffix.

    product_desc: locked canonical product description (Stage 1). FLUX weighs
    the START and END of a prompt most, so the locked product identity is
    stated FIRST and repeated LAST (after the style keywords), and the scene
    palette is explicitly confined to backdrop/props. This pins the product's
    own colors even in bold/colorful styles.
    """
    palette = ", ".join(shot.color_palette[:5])
    suffix = STYLE_SUFFIXES.get(style_preset, "")
    if not product_desc:
        return (f"{shot.image_prompt}. "
                f"Backdrop and props color palette (do not recolor the product): {palette}. {suffix}")
    desc = product_desc.strip().rstrip(".")
    return (f"Commercial product photo of {desc}. "
            f"{shot.image_prompt}. "
            f"Colorful palette for the backdrop, surfaces and props ONLY: {palette} - "
            f"the background may be colorful, the product may not change. {suffix}. "
            f"The product itself stays in its exact original colors: {desc}, "
            f"identical shape, materials and branding, unaltered")

def seed_for(shot_id: int, regen_counter: int = 0) -> int:
    """FR-2.3/FR-2.4: stable per-shot seed; bump regen_counter to reroll one frame."""
    return SEED_BASE + shot_id * 1000 + regen_counter

def generate_frames(shots, style_preset: str, aspect: str = "1:1",
                    regen_counters: dict | None = None,
                    product_desc: str = "") -> list:
    """Render the given shots in ONE backend call. Returns PIL.Images in order.
    regen_counters: {shot_id: int} — bumps the seed of rerolled frames (FR-2.4).
    product_desc: locked product identity prefixed to every prompt."""
    w, h = SIZES.get(aspect, SIZES["1:1"])
    regen_counters = regen_counters or {}
    seeds = [seed_for(s.id, regen_counters.get(s.id, 0)) for s in shots]
    prompts = [build_prompt(s, style_preset, product_desc) for s in shots]
    return flux_generate_batch(prompts, w, h, seeds)

def generate_frame(shot, style_preset: str, aspect: str = "1:1", regen_counter: int = 0,
                   product_desc: str = ""):
    """Render one shot (used by per-frame regen). Returns PIL.Image."""
    return generate_frames([shot], style_preset, aspect,
                           regen_counters={shot.id: regen_counter},
                           product_desc=product_desc)[0]