cli/internal/shot/shot.go

// Package shot captures a PNG screenshot of a rendered CMS page using a
// headless Chromium (chromedp / CDP). It is the image-capture half of the
// `ninja theme screenshot` harness; dojo's Playwright runner only ingests
// pass/fail JSON, so capture lives here.
package shot

import (
	"context"
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/chromedp/cdproto/emulation"
	"github.com/chromedp/cdproto/network"
	"github.com/chromedp/chromedp"
)

// Options configures a single screenshot capture.
type Options struct {
	URL          string        // full page URL to navigate to
	Width        int           // viewport width in px
	Height       int           // viewport height in px
	FullPage     bool          // capture the full scroll height, not just the viewport
	WaitSelector string        // CSS selector to wait for before capturing (e.g. "section")
	Timeout      time.Duration // overall navigation+capture timeout
	Mode         string        // "" (site default), "light" or "dark" — sets the bn-theme cookie + emulated prefers-color-scheme before navigation
	Scale        float64       // device scale factor; 0 = 1.0. 2 captures retina PNGs (Width*2 × Height*2 px)
	Settle       time.Duration // extra wait after WaitSelector before capturing — lets scroll-reveal transitions and canvas animations finish (0 = capture immediately)
}

// PreviewURL composes the gallery-page URL with the render-only override.
// host is the gallery site base (scheme+host), slug is the page path ("/"),
// theme is the bare theme key written to ?preview_template.
func PreviewURL(host, slug, theme string) string {
	host = strings.TrimRight(host, "/")
	if slug == "" {
		slug = "/"
	}
	if !strings.HasPrefix(slug, "/") {
		slug = "/" + slug
	}
	return fmt.Sprintf("%s%s?preview_template=%s", host, slug, url.QueryEscape(theme))
}

// Capture navigates to opts.URL in a headless Chromium and returns PNG bytes.
func Capture(ctx context.Context, opts Options) ([]byte, error) {
	if opts.Width == 0 {
		opts.Width = 1440
	}
	if opts.Height == 0 {
		opts.Height = 900
	}
	if opts.Timeout == 0 {
		opts.Timeout = 30 * time.Second
	}

	allocCtx, cancelAlloc := chromedp.NewExecAllocator(ctx,
		append(chromedp.DefaultExecAllocatorOptions[:],
			chromedp.WindowSize(opts.Width, opts.Height),
			chromedp.Flag("headless", true),
		)...,
	)
	defer cancelAlloc()

	browserCtx, cancelBrowser := chromedp.NewContext(allocCtx)
	defer cancelBrowser()

	timeoutCtx, cancelTimeout := context.WithTimeout(browserCtx, opts.Timeout)
	defer cancelTimeout()

	var buf []byte
	viewport := []chromedp.EmulateViewportOption(nil)
	if opts.Scale > 0 {
		viewport = append(viewport, chromedp.EmulateScale(opts.Scale))
	}
	tasks := chromedp.Tasks{
		chromedp.EmulateViewport(int64(opts.Width), int64(opts.Height), viewport...),
	}
	if opts.Mode == "light" || opts.Mode == "dark" {
		u, err := url.Parse(opts.URL)
		if err != nil {
			return nil, fmt.Errorf("parse %s: %w", opts.URL, err)
		}
		host := u.Hostname()
		mode := opts.Mode
		tasks = append(tasks,
			chromedp.ActionFunc(func(ctx context.Context) error {
				return network.SetCookie("bn-theme", mode).
					WithDomain(host).WithPath("/").Do(ctx)
			}),
			emulation.SetEmulatedMedia().WithFeatures([]*emulation.MediaFeature{
				{Name: "prefers-color-scheme", Value: mode},
			}),
		)
	}
	tasks = append(tasks, chromedp.Navigate(opts.URL))
	if opts.WaitSelector != "" {
		tasks = append(tasks, chromedp.WaitVisible(opts.WaitSelector, chromedp.ByQuery))
	} else {
		tasks = append(tasks, chromedp.WaitReady("body", chromedp.ByQuery))
	}
	if opts.Settle > 0 {
		tasks = append(tasks, chromedp.Sleep(opts.Settle))
	}
	if opts.FullPage {
		tasks = append(tasks, chromedp.FullScreenshot(&buf, 90))
	} else {
		tasks = append(tasks, chromedp.CaptureScreenshot(&buf))
	}

	if err := chromedp.Run(timeoutCtx, tasks); err != nil {
		return nil, fmt.Errorf("capture %s: %w", opts.URL, err)
	}
	return buf, nil
}