pi-weekly-newspaper/src/cover.py

import logging
import os
from collections import Counter
from datetime import date
from io import BytesIO
from urllib.parse import quote

import requests
from PIL import Image as PILImage, ImageDraw, ImageFont

import config

logger = logging.getLogger(__name__)

CATEGORY_PROMPTS = {
    "Government": (
        "Watercolor painting of Plymouth Massachusetts town hall, "
        "New England civic architecture, soft morning light, no text"
    ),
    "Culture": (
        "Plymouth Massachusetts waterfront boardwalk, autumn festival "
        "atmosphere, warm tones, New England charm, no text"
    ),
    "Culture Calendar": (
        "Plymouth Massachusetts harbor at golden hour, sailboats and "
        "historic waterfront, warm autumn palette, no text"
    ),
    "Sports": (
        "Plymouth Massachusetts harbor marina at sunrise, boats and "
        "morning mist on Cape Cod Bay, no text"
    ),
    "Environment": (
        "Cape Cod Bay shoreline near Plymouth Massachusetts, coastal "
        "wetlands and dune grass, golden hour light, no text"
    ),
    "Opinion": (
        "Plymouth Rock historic site in Plymouth Massachusetts, "
        "dramatic sky, iconic New England landmark, no text"
    ),
    "Feature": (
        "Downtown Plymouth Massachusetts street scene, historic "
        "brick storefronts, New England small town charm, no text"
    ),
    "Default": (
        "Aerial view of Plymouth Massachusetts coastal town, serene "
        "harbor, Cape Cod Bay, New England charm, no text"
    ),
}


def _get_font(size: int) -> ImageFont.FreeTypeFont:
    try:
        return ImageFont.truetype(
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", size
        )
    except OSError:
        pass
    try:
        return ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", size)
    except OSError:
        pass
    return ImageFont.load_default()


def _get_dominant_category(categories: list[str]) -> str:
    if not categories:
        return "Default"
    counter = Counter(categories)
    most_common = counter.most_common(1)[0][0]
    return most_common


def _draw_text_overlays(
    draw: ImageDraw.ImageDraw,
    width: int,
    height: int,
    week_start: date,
    week_end: date,
    headlines: list[str],
) -> None:
    title_font = _get_font(38)
    date_font = _get_font(20)
    headline_font = _get_font(14)

    draw.rectangle([(0, 0), (width, 90)], fill=(0, 0, 0, 204))
    draw.text(
        (width // 2, 20), "Plymouth Independent",
        fill="white", font=title_font, anchor="mt",
    )

    iso_week = week_start.isocalendar()[1]
    date_str = (
        f"Week {iso_week} \u00b7 "
        f"{week_start.strftime('%b %d')} \u2013 {week_end.strftime('%b %d, %Y')}"
    )
    draw.text(
        (width // 2, 62), date_str,
        fill=(220, 220, 220), font=date_font, anchor="mt",
    )

    if headlines:
        max_headlines = min(len(headlines), 5)
        strip_height = 28 + max_headlines * 22
        strip_top = height - strip_height
        draw.rectangle(
            [(0, strip_top), (width, height)], fill=(0, 0, 0, 170)
        )
        y = strip_top + 8
        for headline in headlines[:max_headlines]:
            text = f"\u2022 {headline}"
            if len(text) > 45:
                text = text[:42] + "..."
            draw.text((16, y), text, fill="white", font=headline_font)
            y += 22


def generate_programmatic_cover(
    output_dir: str,
    week_start: date,
    week_end: date,
    headlines: list[str],
    categories: list[str],
) -> str:
    os.makedirs(output_dir, exist_ok=True)
    w, h = config.COVER_SIZE
    img = PILImage.new("RGBA", (w, h), color=(245, 245, 245, 255))
    draw = ImageDraw.Draw(img)

    _draw_text_overlays(draw, w, h, week_start, week_end, headlines)

    out = img.convert("RGB")
    filename = f"cover-{week_start.isoformat()}-text.jpg"
    path = os.path.join(output_dir, filename)
    out.save(path, format="JPEG", progressive=False, quality=90)
    return path


def generate_ai_cover(
    output_dir: str,
    week_start: date,
    week_end: date,
    headlines: list[str],
    categories: list[str],
) -> str:
    os.makedirs(output_dir, exist_ok=True)
    w, h = config.COVER_SIZE

    dominant = _get_dominant_category(categories)
    prompt = CATEGORY_PROMPTS.get(dominant, CATEGORY_PROMPTS["Default"])
    url = config.POLLINATIONS_URL.format(prompt=quote(prompt))

    try:
        response = requests.get(url, timeout=60)
        response.raise_for_status()

        bg = PILImage.open(BytesIO(response.content))
        if bg.mode not in ("RGB", "RGBA"):
            bg = bg.convert("RGB")
        bg = bg.resize((w, h), PILImage.Resampling.LANCZOS)

        img = bg.convert("RGBA")
        overlay = PILImage.new("RGBA", (w, h), (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)
        _draw_text_overlays(draw, w, h, week_start, week_end, headlines)
        img = PILImage.alpha_composite(img, overlay)

        out = img.convert("RGB")
        filename = f"cover-{week_start.isoformat()}-ai.jpg"
        path = os.path.join(output_dir, filename)
        out.save(path, format="JPEG", progressive=False, quality=90)
        return path

    except Exception as e:
        logger.error("AI cover generation failed, falling back to programmatic: %s", e)
        return generate_programmatic_cover(
            output_dir, week_start, week_end, headlines, categories
        )


def generate_cover(
    method: str,
    output_dir: str,
    week_start: date,
    week_end: date,
    headlines: list[str],
    categories: list[str],
) -> str:
    if method == "ai":
        return generate_ai_cover(
            output_dir, week_start, week_end, headlines, categories
        )
    return generate_programmatic_cover(
        output_dir, week_start, week_end, headlines, categories
    )