book_builder.py

"""
Book builder — assembles pages into storybook HTML and PDF.

Features:
- Magic Loader: animated wait messages during generation
- Coloring Book: outline pages from same FLUX images
- Styled PDF covers: scrapbook-style cover matching on-screen UI
"""

import io
import logging
import math
import os
import tempfile
import urllib.parse

from fpdf import FPDF
from PIL import Image, ImageDraw, ImageFilter, ImageFont, ImageOps

from config import COLORS

logger = logging.getLogger(__name__)

FONTS_DIR = os.path.join(os.path.dirname(__file__), "assets", "fonts")
FONT_GAEGU = os.path.join(FONTS_DIR, "Gaegu-Bold.ttf")
FONT_CAVEAT = os.path.join(FONTS_DIR, "Caveat.ttf")


# ============================================================================
# STORYBOOK HTML
# ============================================================================

PAGE_HTML = """
<div class="book-page" style="page-break-after: always;">
  <div class="page-art">
    <img src="{img_b64}" alt="Page {page_num}: {alt_text}"/>
  </div>
  <p class="page-text">{text}</p>
  <span class="page-num">~ {page_num} ~</span>
</div>
"""

COVER_HTML = """
<div class="book-cover">
  <p class="cover-kicker">✨ a DoodleBook story ✨</p>
  {cover_art}
  <h2 class="cover-title">{title}</h2>
  {badge}
</div>
"""

ENGINE_BADGES = {
    "flux":   '<span class="engine-badge flux">illustrated by FLUX.2-klein</span>',
    "sketch": '<span class="engine-badge sketch">preview sketch — connect a GPU for FLUX art</span>',
}


def build_book_html(
    pages_images: list,
    pages_texts: list,
    title: str,
    engine: str = "flux",
) -> str:
    """Build storybook HTML from images and texts."""
    badge = ENGINE_BADGES.get(engine, "")
    cover_art = ""
    if pages_images:
        cover_art = (
            f'<div class="cover-art">'
            f'<img src="{_img_src_for_bytes(pages_images[0])}" alt="Cover illustration"/>'
            f'</div>'
        )
    cover = COVER_HTML.format(title=title, badge=badge, cover_art=cover_art)

    pages_html = ""
    for i, (img_bytes, text) in enumerate(zip(pages_images, pages_texts)):
        pages_html += PAGE_HTML.format(
            img_b64=_img_src_for_bytes(img_bytes),
            text=text,
            page_num=i + 1,
            alt_text=(text[:50] + "...") if len(text) > 50 else text,
        )

    return f"""<div class="book-container">
    {cover}
    {pages_html}
    </div>"""


# ============================================================================
# MAGIC LOADER (Feature 1)
# ============================================================================

def magic_loader_html(stage: str = "story", hero_name: str = "your hero") -> str:
    """
    Animated crayon-styled wait panel with rotating messages.

    Pure CSS fade cycle — no JS, no Gradio streaming needed.
    stage: "story" | "images" | "tts" — controls which messages show.
    """
    messages = {
        "story": [
            f"✏️ MiniCPM is dreaming up {hero_name}'s story\u2026",
            "📖 Writing page by page\u2026",
            "💡 Did you know? Your whole storybook runs on tiny models!",
        ],
        "images": [
            f"🎨 FLUX is painting {hero_name}'s 6 pages\u2026",
            "🖌️ Adding crayon details\u2026",
            "💡 Did you know? The brain is only 3B parameters!",
        ],
        "tts": [
            "🔊 VoxCPM is recording the narration\u2026",
            "🎙️ Warming up the storyteller voice\u2026",
            "💡 Did you know? The voice runs on a 2B model!",
        ],
    }
    msgs = messages.get(stage, messages["story"])
    n = len(msgs)
    step = 3.0
    total = step * n

    items = ""
    for i, m in enumerate(msgs):
        delay = i * step
        items += (
            f'<div class="ml-msg" style="animation-delay:{delay:.1f}s;">{m}</div>\n'
        )

    return f"""
    <div class="magic-loader">
      <div class="ml-spinner"></div>
      <div class="ml-stack">
        {items}
      </div>
      <style>
        .magic-loader {{
          text-align:center; padding:36px 20px; font-family:'Gaegu',cursive;
          font-size:22px; color:#2e2a26; position:relative; min-height:120px;
        }}
        .ml-spinner {{
          width:36px; height:36px; margin:0 auto 14px;
          border:4px solid #efe0c2; border-top-color:#ef6a3a;
          border-radius:50%; animation:ml-spin 1s linear infinite;
        }}
        @keyframes ml-spin {{ to {{ transform:rotate(360deg); }} }}
        .ml-stack {{ position:relative; height:64px; }}
        .ml-msg {{
          position:absolute; inset:0; display:flex; align-items:center;
          justify-content:center; text-align:center; padding:0 10px;
          opacity:0; animation:ml-cycle {total:.1f}s infinite;
        }}
        @keyframes ml-cycle {{
          0%,24% {{ opacity:1; }}
          33%,100% {{ opacity:0; }}
        }}
        @media (prefers-reduced-motion:reduce) {{
          .ml-spinner {{ animation:none; }}
          .ml-msg {{ animation:none; opacity:1; position:static; margin:6px 0; }}
        }}
      </style>
    </div>
    """


# ============================================================================
# STYLIZED PDF COVER (Feature 3)
# ============================================================================

def _load_font(path: str, size: int):
    """Load TrueType font with fallback to default."""
    try:
        return ImageFont.truetype(path, size)
    except Exception:
        try:
            return ImageFont.truetype("arial.ttf", size)
        except Exception:
            return ImageFont.load_default()


def _paper_grain(img, rng_seed=42, amount=2000):
    """Sprinkle subtle specks so cover reads as paper."""
    import random
    rng = random.Random(rng_seed)
    px = img.load()
    w, h = img.size
    for _ in range(amount):
        x = rng.randint(0, w - 1)
        y = rng.randint(0, h - 1)
        r, g, b = px[x, y]
        s = rng.choice([-10, -6, 6, 10])
        px[x, y] = (
            max(0, min(255, r + s)),
            max(0, min(255, g + s)),
            max(0, min(255, b + s)),
        )


def _wrap_title(title: str, font, max_w: int) -> list:
    words = title.split()
    lines, current = [], ""
    for w in words:
        test = f"{current} {w}".strip()
        if font.getbbox(test)[2] - font.getbbox(test)[0] > max_w:
            if current:
                lines.append(current)
            current = w
        else:
            current = test
    if current:
        lines.append(current)
    return lines or [title]


def render_cover_image(
    title: str,
    badge_text: str = "illustrated by FLUX.2-klein",
    kind: str = "story",
    first_page_bytes: bytes = None,
) -> bytes:
    """
    Render a full-page illustrated cover as PNG bytes (A4 @150dpi, 1240x1754).

    When first_page_bytes is provided (the FLUX page 1 illustration), it fills
    the cover as a full-bleed background with dark banner overlays for the title.
    Falls back to cream-background design when no illustration is available.
    """
    W, H = 1240, 1754
    sun   = (244, 198, 74)
    leaf  = (116, 184, 90)
    white = (255, 255, 255)
    ink   = (46, 42, 38)
    berry = (214, 81, 122)

    if first_page_bytes:
        # --- Cream card cover: same layout as the in-app HTML cover ---
        cream = (255, 248, 230)
        img = Image.new("RGB", (W, H), cream)
        draw = ImageDraw.Draw(img)
        _paper_grain(img, amount=5000)

        # Kicker
        font_kicker = _load_font(FONT_CAVEAT, 66)
        draw.text((W // 2, 118), "✨  a DoodleBook story  ✨",
                  font=font_kicker, fill=berry, anchor="mm")

        # FLUX illustration — square-crop & resize to 900 px
        src = Image.open(io.BytesIO(first_page_bytes)).convert("RGB")
        min_d = min(src.width, src.height)
        src = src.crop(((src.width - min_d) // 2, (src.height - min_d) // 2,
                        (src.width + min_d) // 2, (src.height + min_d) // 2))
        src = src.resize((900, 900), Image.LANCZOS)

        # Frame: sun inner (12 px) + ink outer (16 px), rounded corners
        pad_inner, pad_outer = 12, 16
        pad = pad_inner + pad_outer
        ix, iy = (W - 900) // 2, 200          # top-left of image
        # Ink outer border
        draw.rounded_rectangle(
            [ix - pad, iy - pad, ix + 900 + pad, iy + 900 + pad],
            radius=30, fill=ink,
        )
        # Sun inner border
        draw.rounded_rectangle(
            [ix - pad_inner, iy - pad_inner, ix + 900 + pad_inner, iy + 900 + pad_inner],
            radius=20, fill=sun,
        )
        # Paste illustration
        img.paste(src, (ix, iy))
        draw = ImageDraw.Draw(img)

        # Title — below frame
        font_title = _load_font(FONT_GAEGU, 100)
        lines = _wrap_title(title, font_title, W - 160)
        line_h = 118
        y0 = iy + 900 + pad + 56
        for i, line in enumerate(lines):
            cy = y0 + i * line_h
            draw.text((W // 2 + 3, cy + 3), line, font=font_title,
                      fill=sun, anchor="mm")
            draw.text((W // 2, cy), line, font=font_title,
                      fill=ink, anchor="mm")
        total_h = len(lines) * line_h

        # Badge
        font_badge = _load_font(FONT_GAEGU, 46)
        bb = font_badge.getbbox(badge_text)
        tw, th = bb[2] - bb[0], bb[3] - bb[1]
        bx = W // 2 - (tw + 44) // 2
        by = max(y0 + total_h + 60, H - 160)
        draw.rounded_rectangle([bx, by, bx + tw + 44, by + th + 24],
                               radius=14, fill=leaf)
        draw.text((W // 2, by + th // 2 + 12), badge_text,
                  font=font_badge, fill=white, anchor="mm")

    else:
        # --- Cream fallback cover ---
        cream = (255, 248, 230)
        img = Image.new("RGB", (W, H), cream)
        draw = ImageDraw.Draw(img)
        _paper_grain(img, amount=6000)

        font_kicker = _load_font(FONT_CAVEAT, 60)
        draw.text((W // 2, 520), "✨ a DoodleBook story ✨",
                  font=font_kicker, fill=berry, anchor="mm")

        font_title = _load_font(FONT_GAEGU, 120)
        lines = _wrap_title(title, font_title, W - 200)
        total_h = len(lines) * 140
        y0 = 680
        for i, line in enumerate(lines):
            cy = y0 + i * 140
            draw.text((W // 2 + 4, cy + 4), line, font=font_title,
                      fill=sun, anchor="mm")
            draw.text((W // 2, cy), line, font=font_title,
                      fill=ink, anchor="mm")

        font_badge = _load_font(FONT_GAEGU, 44)
        bb = font_badge.getbbox(badge_text)
        tw, th = bb[2] - bb[0], bb[3] - bb[1]
        bx = W // 2 - (tw + 40) // 2
        by = y0 + total_h + 80
        draw.rounded_rectangle([bx, by, bx + tw + 40, by + th + 24],
                               radius=16, fill=leaf)
        draw.text((bx + 20 + tw // 2, by + 12 + th // 2), badge_text,
                  font=font_badge, fill=white, anchor="mm")

        import random
        rng = random.Random(77)
        sx, sy = W // 2 - 180, y0 + total_h + 40
        pts = [(sx + j * 10, sy + rng.randint(-5, 5)) for j in range(36)]
        draw.line(pts, fill=leaf, width=5, joint="curve")

    _paper_grain(img, amount=2000)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


# ============================================================================
# COLORING BOOK (Feature 2)
# ============================================================================

COLORING_COVER_HTML = """
<div class="book-cover coloring-cover">
  <p class="cover-kicker">a DoodleBook coloring book</p>
  <h2 class="cover-title">{title}</h2>
  <span class="engine-badge sketch">a coloring book to color in</span>
</div>
"""

COLORING_PAGE_HTML = """
<div class="book-page" style="page-break-after: always;">
  <div class="page-art coloring-art">
    <img src="{img_b64}" alt="Coloring page {page_num}"/>
  </div>
  <p class="page-text coloring-text">{text}</p>
  <span class="page-num">~ {page_num} ~</span>
</div>
"""


def build_coloring_html(
    outline_imgs: list,
    page_texts: list,
    title: str,
) -> str:
    """Build coloring-book HTML — same layout, outline images."""
    cover = COLORING_COVER_HTML.format(title=title)

    pages_html = ""
    for i, (img_bytes, text) in enumerate(zip(outline_imgs, page_texts)):
        pages_html += COLORING_PAGE_HTML.format(
            img_b64=_img_src_for_bytes(img_bytes),
            text=text,
            page_num=i + 1,
        )

    return f"""<div class="book-container">
    {cover}
    {pages_html}
    </div>"""


def _img_src_for_bytes(img_bytes: bytes) -> str:
    """Persist an image to temp and return a Gradio-served file URL.

    Inline base64 made the final HTML payload >10MB for a 6-page book, which
    caused the frontend to stall even though backend generation succeeded.
    File-backed URLs keep the final update small and let the browser fetch page
    images incrementally.
    """
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        tmp.write(img_bytes)
        path = tmp.name
    return f"/gradio_api/file={urllib.parse.quote(path)}"


# ============================================================================
# PDF EXPORT
# ============================================================================

def _pdf_cover_page(pdf: FPDF, cover_png_bytes: bytes):
    """Place a styled PNG cover as the current PDF page."""
    img = Image.open(io.BytesIO(cover_png_bytes))
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        img.save(tmp.name, "PNG")
        pdf.image(tmp.name, x=0, y=0, w=210, h=297)
    os.unlink(tmp.name)


def _safe_text(s: str) -> str:
    """Make text safe for fpdf's Latin-1 core fonts.

    The page text comes from the story generator and can contain Unicode
    punctuation (em-dash —, curly quotes, ellipsis). fpdf's Helvetica is
    Latin-1 only and RAISES on those, which previously killed PDF generation
    and left the download button missing. Map the common ones, then drop any
    remaining non-Latin-1 char so a PDF is ALWAYS produced.
    """
    repl = {
        "—": "-", "–": "-", "‒": "-",
        "‘": "'", "’": "'", "“": '"', "”": '"',
        "…": "...", " ": " ", "​": "",
    }
    for k, v in repl.items():
        s = s.replace(k, v)
    return s.encode("latin-1", "replace").decode("latin-1")


def export_pdf(
    pages_images: list,
    pages_texts: list,
    title: str,
    output_path: str = None,
) -> str:
    """Export storybook as PDF with styled cover."""
    if output_path is None:
        output_path = tempfile.mktemp(suffix=".pdf", prefix="doodlebook_")

    pdf = FPDF()
    pdf.set_auto_page_break(auto=False)

    # Cover: full-bleed FLUX illustration when available
    cover_png = render_cover_image(
        title, "illustrated by FLUX.2-klein", "story",
        first_page_bytes=pages_images[0] if pages_images else None,
    )
    pdf.add_page()
    _pdf_cover_page(pdf, cover_png)

    for i, (img_bytes, text) in enumerate(zip(pages_images, pages_texts)):
        pdf.add_page()
        img = Image.open(io.BytesIO(img_bytes))
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            img.save(tmp.name, "PNG")
            pdf.image(tmp.name, x=10, y=10, w=190)
        os.unlink(tmp.name)

        pdf.set_y(210)
        pdf.set_font("Helvetica", "", 16)
        pdf.multi_cell(0, 10, _safe_text(text), align="C")

        pdf.set_y(270)
        pdf.set_font("Helvetica", "I", 10)
        pdf.cell(0, 10, f"Page {i+1}", ln=True, align="C")

    pdf.output(output_path)
    return output_path


def export_coloring_pdf(
    outline_imgs: list,
    page_texts: list,
    title: str,
    output_path: str = None,
) -> str:
    """Export coloring book as PDF with styled cover."""
    if output_path is None:
        output_path = tempfile.mktemp(suffix=".pdf", prefix="doodlebook_coloring_")

    pdf = FPDF()
    pdf.set_auto_page_break(auto=False)

    # Cover: first coloring-page illustration as full-bleed background
    cover_png = render_cover_image(
        title, "a coloring book to color in", "coloring",
        first_page_bytes=outline_imgs[0] if outline_imgs else None,
    )
    pdf.add_page()
    _pdf_cover_page(pdf, cover_png)

    for i, (img_bytes, text) in enumerate(zip(outline_imgs, page_texts)):
        pdf.add_page()
        img = Image.open(io.BytesIO(img_bytes))
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            img.save(tmp.name, "PNG")
            pdf.image(tmp.name, x=10, y=10, w=190)
        os.unlink(tmp.name)

        pdf.set_y(210)
        pdf.set_font("Helvetica", "", 14)
        pdf.multi_cell(0, 10, _safe_text(text), align="C")

        pdf.set_y(270)
        pdf.set_font("Helvetica", "I", 10)
        pdf.cell(0, 10, f"Page {i+1}", ln=True, align="C")

    pdf.output(output_path)
    return output_path