adam/showcase.py

from __future__ import annotations

import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from adam.executor import ToolContext, ToolExecutionError
from adam.models import ExecutionPlan, PlanStep


def build_showcase_plan(
    generation_plans: list[ExecutionPlan],
    *,
    title: str,
    display_seconds: int,
    resolution: str,
    model_settings: list[dict[str, Any]],
) -> ExecutionPlan:
    """Append a durable MP4 render to sequential DDPM/Flow generation plans."""
    usable = [plan for plan in generation_plans if plan.steps]
    if not usable:
        raise ValueError("A showcase needs at least one model.")
    seconds = int(display_seconds)
    if seconds not in {3, 4, 5}:
        raise ValueError("Showcase images must last 3, 4, or 5 seconds.")
    render_step = PlanStep(
        tool_id="showcase_video_renderer",
        title="Render showcase video",
        description="Compose the generated images and animated request interface into an MP4.",
        arguments={
            "title": title.strip() or "ADAM Generation Showcase",
            "display_seconds": seconds,
            "resolution": resolution,
            "models": model_settings,
        },
    )
    image_total = sum(
        int(plan.steps[0].arguments.get("image_count", 0) or 0) for plan in usable
    )
    reasons = [plan.confirmation_reason for plan in usable if plan.confirmation_reason]
    return ExecutionPlan(
        request=f"Create a finished showcase video with {len(usable)} models.",
        summary=(
            f"Generate {image_total} images with {len(usable)} DDPM/Flow models, "
            f"show each for {seconds} seconds, then export an MP4."
        ),
        steps=[step for plan in usable for step in plan.steps] + [render_step],
        requires_confirmation=any(plan.requires_confirmation for plan in usable),
        confirmation_reason="; ".join(dict.fromkeys(reasons)),
        project_name="Showcase Video",
    )


def _safe_filename(value: str) -> str:
    value = re.sub(r'[<>:"/\\|?*\x00-\x1f]+', " ", value.strip())
    return re.sub(r"\s+", " ", value).strip(" .")[:80] or "ADAM Showcase"


def _font(size: int, *, bold: bool = False):
    from PIL import ImageFont

    names = [
        "C:/Windows/Fonts/seguisb.ttf" if bold else "C:/Windows/Fonts/segoeui.ttf",
        "C:/Windows/Fonts/arialbd.ttf" if bold else "C:/Windows/Fonts/arial.ttf",
    ]
    for name in names:
        try:
            return ImageFont.truetype(name, size)
        except OSError:
            pass
    return ImageFont.load_default()


def _fit_image(path: Path, size: tuple[int, int]):
    from PIL import Image, ImageEnhance

    target_w, target_h = size
    with Image.open(path) as source:
        source = source.convert("RGB")
        # Keep the complete generated image visible while enlarging it to the
        # available preview area. The caller provides the 10px frame inset.
        fit_scale = min(target_w / source.width, target_h / source.height)
        fitted_size = (
            max(1, round(source.width * fit_scale)),
            max(1, round(source.height * fit_scale)),
        )
        source = source.resize(fitted_size, Image.Resampling.LANCZOS)
        canvas = Image.new("RGB", size, (1, 7, 25))
        canvas.paste(source, ((target_w - source.width) // 2, (target_h - source.height) // 2))
        source = ImageEnhance.Color(source).enhance(1.18)
        source = ImageEnhance.Contrast(source).enhance(1.07)
        canvas.paste(source, ((target_w - source.width) // 2, (target_h - source.height) // 2))
        return canvas


def _shorten(draw, text: str, font, width: int) -> str:
    if draw.textbbox((0, 0), text, font=font)[2] <= width:
        return text
    value = text
    while value and draw.textbbox((0, 0), value + "…", font=font)[2] > width:
        value = value[:-1]
    return value.rstrip() + "…"


def _compose_frame(
    image_path: Path,
    *,
    title: str,
    models: list[dict[str, Any]],
    model_index: int,
    image_index: int,
    image_count: int,
    size: tuple[int, int],
):
    from PIL import Image, ImageDraw

    width, height = size
    scale = width / 1920
    canvas = Image.new("RGB", size, (1, 5, 18))
    draw = ImageDraw.Draw(canvas)
    # Subtle broadcast-style bands keep the frame readable without relying on assets.
    for y in range(height):
        blue = int(30 + 42 * (1 - y / max(1, height)))
        draw.line((0, y, width, y), fill=(1, 4 + blue // 7, blue))

    margin = int(24 * scale)
    header_h = int(150 * scale)
    footer_h = int(205 * scale)
    panel_w = int(430 * scale)
    gap = int(24 * scale)
    cyan, white, muted, yellow = (20, 222, 255), (244, 249, 255), (142, 172, 213), (255, 202, 20)
    border = (30, 102, 255)
    panel_fill, image_fill = (3, 14, 42), (1, 7, 25)
    draw.rounded_rectangle((margin, margin, width - margin, header_h), radius=int(18 * scale), fill=panel_fill, outline=border, width=max(2, int(3 * scale)))
    kicker = _font(max(14, int(25 * scale)), bold=True)
    heading = _font(max(28, int(68 * scale)), bold=True)
    body = _font(max(14, int(25 * scale)))
    small = _font(max(12, int(20 * scale)))
    request_font = _font(max(13, int(23 * scale)), bold=True)
    draw.text((margin + int(28 * scale), margin + int(18 * scale)), "ADAM GENERATION SERIES", font=kicker, fill=cyan)
    draw.text((margin + int(28 * scale), margin + int(49 * scale)), _shorten(draw, title.upper(), heading, width - int(330 * scale)), font=heading, fill=white)
    draw.text(
        (width - margin - int(28 * scale), margin + int(48 * scale)),
        "FINISHED SHOWCASE", font=kicker, fill=(255, 74, 112), anchor="ra",
    )

    content_top = header_h + gap
    content_bottom = height - footer_h - margin
    draw.rounded_rectangle((margin, content_top, panel_w, content_bottom), radius=int(16 * scale), fill=panel_fill, outline=border, width=max(2, int(2 * scale)))
    draw.text((margin + int(22 * scale), content_top + int(20 * scale)), "REQUEST LIST", font=kicker, fill=cyan)

    row_h = max(30, int(45 * scale))
    list_top = content_top + int(64 * scale)
    visible = max(1, int((content_bottom - list_top - int(20 * scale)) / row_h))
    start = max(0, min(model_index - visible // 2, len(models) - visible))
    end = min(len(models), start + visible)
    for visible_row, idx in enumerate(range(start, end)):
        y = list_top + visible_row * row_h
        active = idx == model_index
        if active:
            draw.rounded_rectangle((margin + int(12 * scale), y, panel_w - int(12 * scale), y + row_h - int(5 * scale)), radius=int(8 * scale), fill=(4, 84, 164), outline=cyan, width=max(1, int(2 * scale)))
        number = f"{idx + 1}."
        draw.text((margin + int(22 * scale), y + int(7 * scale)), number, font=request_font, fill=cyan)
        name = _shorten(draw, str(models[idx].get("name", "Model")), request_font, panel_w - margin - int(95 * scale))
        draw.text((margin + int(72 * scale), y + int(7 * scale)), name, font=request_font, fill=white if active else muted)

    image_left = panel_w + gap
    image_right = width - margin
    image_bottom = content_bottom
    draw.rounded_rectangle((image_left, content_top, image_right, image_bottom), radius=int(16 * scale), fill=image_fill, outline=(139, 46, 255), width=max(2, int(3 * scale)))
    inset = max(10, int(10 * scale))
    fitted = _fit_image(image_path, (image_right - image_left - inset * 2, image_bottom - content_top - inset * 2))
    canvas.paste(fitted, (image_left + inset, content_top + inset))

    current = models[model_index]
    footer_top = height - footer_h
    draw.rounded_rectangle((margin, footer_top, width - margin, height - margin), radius=int(16 * scale), fill=panel_fill, outline=border, width=max(2, int(2 * scale)))
    footer_label = _font(max(11, int(19 * scale)), bold=True)
    footer_value = _font(max(21, int(42 * scale)), bold=True)
    footer_minor = _font(max(11, int(18 * scale)), bold=True)
    footer_y = footer_top + int(23 * scale)
    x = margin + int(28 * scale)
    draw.text((x, footer_y), "CURRENT REQUEST", font=footer_label, fill=muted)
    draw.text((x, footer_y + int(27 * scale)), _shorten(draw, str(current.get("name", "Model")).upper(), footer_value, int(640 * scale)), font=footer_value, fill=white)
    trainer = str(current.get("trainer_label", current.get("trainer", "MODEL"))).upper()
    draw.text((x, footer_y + int(75 * scale)), trainer, font=footer_minor, fill=cyan)
    x2 = int(820 * scale)
    draw.text((x2, footer_y), "IMAGE", font=footer_label, fill=muted)
    draw.text((x2, footer_y + int(27 * scale)), f"{image_index + 1} / {image_count}", font=footer_value, fill=yellow)
    x3 = int(1180 * scale)
    draw.text((x3, footer_top + int(18 * scale)), f"{current.get('steps', '—')} STEPS", font=kicker, fill=white)
    draw.text((x3, footer_top + int(58 * scale)), str(current.get("sampler", "")), font=footer_value, fill=cyan)
    draw.text((x3, footer_top + int(122 * scale)), str(current.get("aspect_ratio", "")), font=footer_minor, fill=muted)
    return canvas


def render_showcase_video(
    context: ToolContext,
    title: str,
    display_seconds: int,
    resolution: str,
    models: list[dict[str, Any]],
) -> dict[str, object]:
    """Render images generated earlier in this same job into a showcase MP4."""
    try:
        import cv2
        import numpy as np
    except ImportError as exc:
        raise ToolExecutionError("Showcase export requires OpenCV and NumPy.") from exc
    if not isinstance(models, list) or not models:
        raise ToolExecutionError("The showcase has no selected models.")
    seconds = int(display_seconds)
    if seconds not in {3, 4, 5}:
        raise ToolExecutionError("Image duration must be 3, 4, or 5 seconds.")
    sizes = {"720p": (1280, 720), "1080p": (1920, 1080)}
    if resolution not in sizes:
        raise ToolExecutionError("Showcase resolution must be 720p or 1080p.")

    history_root = context.root / "data" / "generations"
    records: list[dict[str, Any]] = []
    if history_root.is_dir():
        for metadata in history_root.rglob(f"*{context.job_id}*.json"):
            try:
                payload = json.loads(metadata.read_text(encoding="utf-8"))
            except (OSError, ValueError, TypeError, json.JSONDecodeError):
                continue
            images = [Path(str(item)) for item in payload.get("images", [])]
            images = [item for item in images if item.is_file()]
            if images:
                records.append({
                    "model_name": str(payload.get("model_name", "")),
                    "model_path": str(payload.get("model_path", "")),
                    "images": images,
                })
    record_by_name = {record["model_name"]: record for record in records}
    record_by_path = {record["model_path"]: record for record in records if record["model_path"]}
    slides: list[tuple[Path, int, int, int]] = []
    for model_index, model in enumerate(models):
        record = record_by_path.get(str(model.get("path", ""))) or record_by_name.get(
            str(model.get("name", ""))
        )
        if not record:
            continue
        count = len(record["images"])
        slides.extend((path, model_index, image_index, count) for image_index, path in enumerate(record["images"]))
    if not slides:
        raise ToolExecutionError("No generated showcase images were found for this job.")

    output = context.root / "data" / "showcase_videos"
    output.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    destination = output / f"{_safe_filename(title)}_{timestamp}_{context.job_id}.mp4"
    width, height = sizes[resolution]
    fps = 24
    writer = cv2.VideoWriter(str(destination), cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
    if not writer.isOpened():
        raise ToolExecutionError("Could not open the MP4 video encoder.")
    try:
        frames_per_slide = seconds * fps
        for slide_index, (path, model_index, image_index, image_count) in enumerate(slides):
            context.checkpoint()
            frame_image = _compose_frame(
                path, title=title, models=models, model_index=model_index,
                image_index=image_index, image_count=image_count, size=(width, height),
            )
            frame = cv2.cvtColor(np.asarray(frame_image), cv2.COLOR_RGB2BGR)
            for frame_index in range(frames_per_slide):
                if frame_index % fps == 0:
                    context.checkpoint()
                writer.write(frame)
            context.progress(
                round((slide_index + 1) * 100 / len(slides)),
                f"Rendering showcase image {slide_index + 1} of {len(slides)}",
            )
    finally:
        writer.release()
    if not destination.is_file() or destination.stat().st_size == 0:
        raise ToolExecutionError("The showcase encoder did not produce a video file.")
    manifest = destination.with_suffix(".json")
    manifest.write_text(json.dumps({
        "version": 1, "title": title, "video": str(destination),
        "display_seconds": seconds, "resolution": resolution, "fps": fps,
        "models": models, "image_count": len(slides),
        "created_at": datetime.now(timezone.utc).isoformat(),
    }, indent=2), encoding="utf-8")
    context.log(f"Showcase video saved to {destination}")
    return {
        "output_folder": str(output),
        "assets": [{"kind": "video", "name": title, "path": str(destination), "trainer": "showcase"}],
    }