import hashlib import os import re import sys import urllib.request import time import shutil from contextlib import contextmanager from pathlib import Path from html2image import Html2Image from PIL import Image from email.utils import make_msgid from jinja2 import Environment, FileSystemLoader from loguru import logger from ..config.paths import HTML_PATHS, PATHS from ..config.colors import ROW_COLORS RENDER_HEIGHT = 1080 TRACK_SIZE = (1020, 102) LEGEND_SIZE = (1020, 540) EC2_FLAGS = [ "--no-sandbox", "--disable-dev-shm-usage", "--run-all-compositor-stages-before-draw", ] @contextmanager def silence_stderr(): saved_fd = os.dup(2) devnull_fd = None try: devnull_fd = os.open(os.devnull, os.O_WRONLY) os.dup2(devnull_fd, 2) os.close(devnull_fd) devnull_fd = None yield finally: if devnull_fd is not None: try: os.close(devnull_fd) except OSError: pass os.dup2(saved_fd, 2) os.close(saved_fd) class ScreenshotService: def __init__(self, output_path: Path | None = None, browser_executable: str | None = None): self.output_path = output_path or PATHS.html_tmp kwargs: dict = { "output_path": str(self.output_path), "temp_path": str(self.output_path), "disable_logging": True, } if browser_executable: kwargs["browser_executable"] = browser_executable self.hti = Html2Image(**kwargs) self.hti.browser.use_new_headless = True self.hti.browser.flags += EC2_FLAGS def download_covers_batch(self, urls: list[str]) -> dict[str, str]: """Download a deduplicated set of cover art URLs. Returns {url: local_filename}.""" cache: dict[str, str] = {} unique_urls = set(u for u in urls if u and str(u) != "nan") logger.info(f"Downloading {len(unique_urls)} unique cover images") for url in unique_urls: filename = f"cover_{hashlib.md5(url.encode()).hexdigest()}.png" dest = self.output_path / filename for attempt in range(5): try: urllib.request.urlretrieve(url, dest) cache[url] = filename break except Exception as e: logger.warning(f"Failed to download cover (attempt {attempt + 1}): {e}") time.sleep(5) else: cache[url] = filename return cache def get_cover_path(self, url: str, cover_cache: dict[str, str]) -> str: """Look up a cover from the pre-downloaded cache, or download on cache miss.""" if not url or str(url) == "nan": return "cover_missing.png" cached = cover_cache.get(url) if cached: return cached # Cache miss — download inline as fallback filename = f"cover_{hashlib.md5(url.encode()).hexdigest()}.png" dest = self.output_path / filename for attempt in range(3): try: urllib.request.urlretrieve(url, dest) cover_cache[url] = filename return filename except Exception as e: logger.warning(f"Cover cache miss, download attempt {attempt + 1} failed: {e}") time.sleep(2) return "cover_missing.png" def screenshot_row(self, track_index: int, row_html: str) -> str: filename = f"track_{track_index}.png" with silence_stderr(): self.hti.screenshot( html_str=row_html, save_as=filename, size=(TRACK_SIZE[0], RENDER_HEIGHT), ) self._crop(filename, TRACK_SIZE) return self._make_cid(track_index) def screenshot_multi_market_row(self, track_index: int, row_html: str, n_country_rows: int) -> str: """Screenshot a multi-market card with dynamic height based on number of countries.""" filename = f"track_{track_index}.png" # Card content: header (80px) + table header row (25px) + country rows + bottom padding card_content_height = 80 + 25 + (n_country_rows * 25) + 10 # Inject max-height on the card div to prevent background overflow constrained_html = row_html.replace( 'class="card"', f'class="card" style="height: {card_content_height}px;"', ) # Total image: 4px top padding + card content + 16px bottom margin total_height = 4 + card_content_height + 16 with silence_stderr(): self.hti.screenshot( html_str=constrained_html, save_as=filename, size=(TRACK_SIZE[0], RENDER_HEIGHT), ) self._crop(filename, (TRACK_SIZE[0], total_height)) return self._make_cid(track_index) def screenshot_legend(self, disclaimers: dict[str, str], multi_market: bool = False) -> str: env = Environment( loader=FileSystemLoader(str(HTML_PATHS.legend.parent)), autoescape=False, ) template_path = HTML_PATHS.legend_multi_market if multi_market else HTML_PATHS.legend logger.debug(f"Using legend template: {template_path.name} (multi_market={multi_market})") template = env.get_template(template_path.name) legend_html = template.render( default_color=ROW_COLORS["default"], fresh_color=ROW_COLORS["freshly_cookin"], new_release_color=ROW_COLORS["new_release"], chart_color=ROW_COLORS["charting"], chart_color2=ROW_COLORS["charting_alt"], spotify_disclaimer=disclaimers.get("spotify", ""), tiktok_disclaimer=disclaimers.get("tiktok", ""), meta_disclaimer=disclaimers.get("meta", ""), shazam_disclaimer=disclaimers.get("shazam", ""), ) with silence_stderr(): self.hti.screenshot( html_str=legend_html, save_as="track_legend.png", size=(LEGEND_SIZE[0], RENDER_HEIGHT), ) self._crop("track_legend.png", LEGEND_SIZE) return self._make_cid("legend") def prepare_spotify_logo(self) -> str: src = HTML_PATHS.icons / "spotify.png" dst = self.output_path / "track_spotify_logo.png" shutil.copyfile(src, dst) return self._make_cid("spotify_logo") def _crop(self, filename: str, target: tuple[int, int]) -> None: path = self.output_path / filename img = Image.open(path) img = img.crop((0, 0, target[0], target[1])) img.save(path) def _make_cid(self, identifier: int | str) -> str: cid = make_msgid(domain=f"{self.output_path}/track_{identifier}.png") return re.sub(r"@.*?>", ">", cid) def get_track_image_path(self, index: int) -> Path: return self.output_path / f"track_{index}.png" def get_legend_path(self) -> Path: return self.output_path / "track_legend.png" def get_spotify_logo_path(self) -> Path: return self.output_path / "track_spotify_logo.png" def cleanup(self) -> None: """Remove track screenshot PNGs (per-report). Covers are retained for reuse.""" for filename in os.listdir(self.output_path): file_path = self.output_path / filename try: if file_path.is_file() and file_path.suffix in (".png",): if filename.startswith("track_") and filename != "track_spotify_logo.png": os.unlink(file_path) except OSError as e: logger.warning(f"Failed to delete {file_path}: {e}") def cleanup_covers(self) -> None: """Remove downloaded cover art PNGs. Call once after all reports are sent.""" for filename in os.listdir(self.output_path): file_path = self.output_path / filename try: if file_path.is_file() and filename.startswith("cover_"): os.unlink(file_path) except OSError as e: logger.warning(f"Failed to delete {file_path}: {e}")