"""Stage 3 — keyframes. Two backends:
  hunyuan-frame : text-to-video for N frames, steal frame 0. No extra weights.
  flux / sdxl   : a real image model, resolved from DWCT_IMAGE_MODEL."""
from __future__ import annotations

import logging
from pathlib import Path

from . import config, models

LOG = logging.getLogger("dwct.keyframes")


def _apply_palette(prompt: str, palette: list[str]) -> str:
    if not palette:
        return prompt
    return f"{prompt}, colour palette {' '.join(palette)}"


def _via_hunyuan(prompt: str, negative: str, out: Path, seed: int) -> Path:
    import torch
    from PIL import Image

    pipe = models.hunyuan()
    gen = torch.Generator(device="cpu").manual_seed(seed)
    common = dict(prompt=prompt, negative_prompt=negative,
                  num_inference_steps=config.STEPS, generator=gen,
                  height=config.HEIGHT, width=config.WIDTH)
    try:
        frames = pipe(num_frames=5, **common).frames[0]
    except TypeError:
        frames = pipe(**common).frames[0]
    Image.fromarray(frames[0]).save(out)
    return out


def _via_image_model(prompt: str, negative: str, out: Path, seed: int) -> Path:
    import torch

    pipe = models.image_model()
    if pipe is None:
        raise RuntimeError("DWCT_IMAGE_MODEL is unset")
    gen = torch.Generator(device="cpu").manual_seed(seed)
    img = pipe(prompt=prompt, negative_prompt=negative,
               num_inference_steps=config.STEPS, guidance_scale=3.5,
               height=config.HEIGHT, width=config.WIDTH, generator=gen).images[0]
    img.save(out)
    return out


def make(prompt: str, negative: str, out: Path, seed: int,
         palette: list[str]) -> Path:
    out.parent.mkdir(parents=True, exist_ok=True)
    backend = config.KEYFRAME_BACKEND
    if backend == "auto":
        backend = "flux" if config.IMAGE_MODEL else "hunyuan-frame"

    p = _apply_palette(prompt, palette)
    LOG.info("keyframe -> %s via %s", out.name, backend)
    if backend == "hunyuan-frame":
        return _via_hunyuan(prompt, negative, out, seed)
    return _via_image_model(p, negative, out, seed)