pi-ext

import type { ThemeColor } from "@earendil-works/pi-coding-agent";
import {
	allocateImageId,
	type Component,
	getCapabilities,
	getCellDimensions,
	Image,
} from "@earendil-works/pi-tui";
import type { Browser, Page } from "playwright";
import type { MockupRequest } from "./mockup.js";

const PREVIEW_DURATION_MS = 4_000;
const STATIC_OBSERVATION_MS = 600;
const LOOP_PAUSE_MS = 1_000;
const MAX_REPLAY_FRAME_SETS = 120;
const MAX_REPLAY_BYTES = 16 * 1024 * 1024;
const MAX_SNAPSHOT_BYTES = 16 * 1024 * 1024;
const MAX_REPLAY_LOOPS = 10;
const MAX_PLAYBACK_FPS = 30;
const PREVIEW_WIDTH_PX = 680;
const PREVIEW_HEIGHT_PX = 304;
const MAX_PREVIEW_HEIGHT_PX = 720;
const PREVIEW_BORDER_PX = 4;

export type PreviewFrameSet = {
	capturedAt: number;
	frames: string[];
	cropped?: boolean[];
};

export type MockupPreview = {
	readonly labels: string[];
	readonly done: Promise<void>;
	current(): readonly string[] | undefined;
	cropped(): readonly boolean[];
	expired(): boolean;
	version(): number;
	subscribe(listener: () => void): () => void;
	close(): void;
};

export type MockupPreviewManager = {
	start(
		id: string,
		request: MockupRequest,
		signal?: AbortSignal,
	): Promise<MockupPreview>;
	get(id: string): MockupPreview | undefined;
	close(): Promise<void>;
};

export class PreviewCancelledError extends Error {
	constructor() {
		super("mockup preview cancelled");
		this.name = "PreviewCancelledError";
	}
}

export function createMockupPreviewManager(): MockupPreviewManager {
	const previews = new Map<string, PreviewSession>();
	let browser: Browser | undefined;
	let browserPromise: Promise<Browser> | undefined;
	let generation = 0;
	let managerClosed = false;

	const acquireBrowser = async (): Promise<Browser> => {
		if (managerClosed) throw new PreviewCancelledError();
		if (browser?.isConnected()) return browser;
		const launch =
			browserPromise ??
			import("playwright").then(({ chromium }) =>
				chromium.launch({ headless: true, timeout: 10_000 }),
			);
		browserPromise = launch;
		try {
			const candidate = await launch;
			if (managerClosed) {
				await candidate.close().catch(() => undefined);
				throw new PreviewCancelledError();
			}
			browser = candidate;
			return candidate;
		} finally {
			if (browserPromise === launch) browserPromise = undefined;
		}
	};

	const pruneSnapshots = (): void => {
		let bytes = [...previews.values()].reduce(
			(total, preview) => total + preview.snapshotBytes(),
			0,
		);
		for (const preview of previews.values()) {
			if (bytes <= MAX_SNAPSHOT_BYTES) break;
			bytes -= preview.snapshotBytes();
			preview.close();
		}
	};

	return {
		async start(id, request, signal) {
			const owner = ++generation;
			for (const preview of previews.values()) preview.freeze();
			pruneSnapshots();
			const assertActive = () => {
				if (managerClosed || signal?.aborted || owner !== generation)
					throw new PreviewCancelledError();
			};
			assertActive();
			const activeBrowser = await acquireBrowser();
			assertActive();
			const page = await activeBrowser.newPage({
				viewport: { width: 720, height: 480 },
				deviceScaleFactor: previewDeviceScaleFactor(
					request.options.length,
					process.stdout.columns,
				),
			});
			const abort = () => void page.close().catch(() => undefined);
			let preview: PreviewSession | undefined;
			signal?.addEventListener("abort", abort, { once: true });
			try {
				assertActive();
				page.setDefaultTimeout(5_000);
				page.setDefaultNavigationTimeout(5_000);
				await page.setContent(renderMockupPreviewPage(request), {
					waitUntil: "load",
					timeout: 5_000,
				});
				await page.waitForFunction(
					() => document.documentElement.dataset.mockupPreviewReady === "true",
					undefined,
					{ timeout: 5_000 },
				);
				assertActive();
				preview = new PreviewSession(request, page, () => {
					if (previews.get(id) === preview) previews.delete(id);
				});
				await preview.start();
				assertActive();
				previews.set(id, preview);
				return preview;
			} catch (error) {
				preview?.close();
				await page.close().catch(() => undefined);
				if (signal?.aborted || owner !== generation || managerClosed)
					throw new PreviewCancelledError();
				throw error;
			} finally {
				signal?.removeEventListener("abort", abort);
			}
		},
		get: (id) => previews.get(id),
		async close() {
			managerClosed = true;
			generation++;
			for (const preview of previews.values()) preview.close();
			previews.clear();
			const launching = browserPromise;
			if (launching)
				await launching.then((value) => value.close()).catch(() => undefined);
			await browser?.close().catch(() => undefined);
			browser = undefined;
		},
	};
}

class PreviewSession implements MockupPreview {
	readonly labels: string[];
	readonly done: Promise<void>;
	private readonly listeners = new Set<() => void>();
	private readonly resolveDone: () => void;
	private readonly rejectDone: (error: Error) => void;
	private recording: PreviewFrameSet[] = [];
	private recordingDurationMs = 0;
	private currentFrames: string[] | undefined;
	private currentCropped: boolean[] = [];
	private currentVersion = 0;
	private completedReplayLoops = 0;
	private disposed = false;
	private frozen = false;
	private doneSettled = false;
	private replayTimer: NodeJS.Timeout | undefined;

	constructor(
		private readonly request: MockupRequest,
		private readonly page: Page,
		private readonly onClose: () => void,
	) {
		this.labels = request.options.map((option) => option.label);
		let resolveDone!: () => void;
		let rejectDone!: (error: Error) => void;
		this.done = new Promise<void>((resolve, reject) => {
			resolveDone = resolve;
			rejectDone = reject;
		});
		this.resolveDone = resolveDone;
		this.rejectDone = rejectDone;
	}

	async start(): Promise<void> {
		const first = await this.capture();
		this.present(first);
		void this.produce(first).then(
			() => this.finish(),
			(error) =>
				this.finish(error instanceof Error ? error : new Error(String(error))),
		);
	}

	current(): readonly string[] | undefined {
		return this.currentFrames;
	}

	cropped(): readonly boolean[] {
		return this.currentCropped;
	}

	expired(): boolean {
		return this.disposed;
	}

	version(): number {
		return this.currentVersion;
	}

	subscribe(listener: () => void): () => void {
		this.listeners.clear();
		this.listeners.add(listener);
		return () => this.listeners.delete(listener);
	}

	snapshotBytes(): number {
		return this.frozen && this.currentFrames
			? this.currentFrames.reduce((total, image) => total + image.length, 0)
			: 0;
	}

	freeze(): void {
		if (this.disposed || this.frozen) return;
		this.frozen = true;
		if (this.replayTimer) clearTimeout(this.replayTimer);
		this.replayTimer = undefined;
		const final = this.recording.at(-1);
		if (final) this.present(final);
		this.recording = [];
		void this.page.close().catch(() => undefined);
		if (!this.doneSettled) {
			this.doneSettled = true;
			this.resolveDone();
		}
	}

	close(): void {
		if (this.disposed) return;
		this.freeze();
		this.disposed = true;
		this.currentFrames = undefined;
		this.currentVersion++;
		for (const listener of this.listeners) listener();
		this.listeners.clear();
		void this.page.close().catch(() => undefined);
		if (!this.doneSettled) {
			this.doneSettled = true;
			this.resolveDone();
		}
		this.onClose();
	}

	private async produce(first: PreviewFrameSet): Promise<void> {
		const startedAt = performance.now();
		this.retain(first);
		let previous = first;
		let sawChange = false;
		while (
			!this.disposed &&
			!this.frozen &&
			performance.now() - startedAt < PREVIEW_DURATION_MS
		) {
			const untilNextFrame =
				1000 / MAX_PLAYBACK_FPS - (performance.now() - previous.capturedAt);
			if (untilNextFrame > 0)
				await new Promise((resolve) => setTimeout(resolve, untilNextFrame));
			const next = await this.capture();
			if (this.disposed || this.frozen) return;
			if (
				!sameFrames(previous.frames, next.frames) ||
				!sameFlags(previous.cropped, next.cropped)
			) {
				sawChange = true;
				this.retain(next);
				this.present(next);
			}
			previous = next;
			if (!sawChange && performance.now() - startedAt >= STATIC_OBSERVATION_MS)
				break;
		}
		this.recordingDurationMs = previous.capturedAt - first.capturedAt;
		if (
			!this.disposed &&
			!this.frozen &&
			sawChange &&
			this.recording.length > 1
		)
			this.scheduleReplay(0, this.endPauseMs());
	}

	private async capture(): Promise<PreviewFrameSet> {
		const cards = this.page.locator(".preview-card");
		const frames: string[] = [];
		for (let index = 0; index < this.request.options.length; index++) {
			const png = await cards.nth(index).screenshot({
				type: "png",
				animations: "allow",
				omitBackground: true,
				timeout: 5_000,
			});
			frames.push(png.toString("base64"));
			if (
				frames.reduce((total, frame) => total + frame.length, 0) >
				MAX_REPLAY_BYTES
			)
				throw new Error("Mockup preview frame set exceeds the 16 MiB limit");
		}
		const cropped = await cards.evaluateAll((elements) =>
			elements.map((card) => {
				const content = card.querySelector<HTMLElement>(".fragment-content");
				if (!content) return false;
				const value =
					content.scrollHeight > content.clientHeight ||
					content.scrollWidth > content.clientWidth;
				card.toggleAttribute("data-preview-cropped", value);
				return value;
			}),
		);
		return { capturedAt: performance.now(), frames, cropped };
	}

	private retain(frame: PreviewFrameSet): void {
		this.recording = thinReplayFrames([...this.recording, frame]);
	}

	private scheduleReplay(index: number, delayMs: number): void {
		this.replayTimer = setTimeout(() => {
			this.replayTimer = undefined;
			if (this.disposed || this.frozen) return;
			const frame = this.recording[index];
			if (!frame) return;
			this.present(frame);
			const nextIndex = index + 1;
			if (nextIndex >= this.recording.length) {
				this.completedReplayLoops++;
				if (this.completedReplayLoops < MAX_REPLAY_LOOPS)
					this.scheduleReplay(0, this.endPauseMs());
				return;
			}
			const next = this.recording[nextIndex];
			this.scheduleReplay(
				nextIndex,
				Math.max(1000 / MAX_PLAYBACK_FPS, next.capturedAt - frame.capturedAt),
			);
		}, delayMs);
		this.replayTimer.unref?.();
	}

	private endPauseMs(): number {
		const final = this.recording.at(-1);
		const first = this.recording[0];
		const uncapturedTail =
			first && final
				? Math.max(
						0,
						this.recordingDurationMs - (final.capturedAt - first.capturedAt),
					)
				: 0;
		return uncapturedTail + LOOP_PAUSE_MS;
	}

	private present(frame: PreviewFrameSet): void {
		this.currentFrames = frame.frames;
		this.currentCropped = frame.cropped ?? [];
		this.currentVersion++;
		for (const listener of this.listeners) listener();
	}

	private finish(error?: Error): void {
		if (this.disposed || this.doneSettled) return;
		this.doneSettled = true;
		void this.page.close().catch(() => undefined);
		if (error) this.rejectDone(error);
		else this.resolveDone();
	}
}

export class MockupPreviewComponent implements Component {
	private imageIds: number[];
	private readonly unsubscribe: () => void;
	private cachedLines: string[] | undefined;
	private complete = false;
	private cachedVersion = -1;
	private cachedWidth = -1;

	constructor(
		private readonly preview: MockupPreview,
		private color: (style: ThemeColor, text: string) => string,
		private requestRender: () => void,
		private showImages: boolean,
	) {
		this.imageIds = preview.labels.map(() => allocateImageId());
		this.unsubscribe = preview.subscribe(() => {
			this.invalidate();
			this.requestRender();
		});
	}

	update(
		color: (style: ThemeColor, text: string) => string,
		requestRender: () => void,
		showImages: boolean,
	): void {
		this.color = color;
		this.requestRender = requestRender;
		if (this.showImages !== showImages) {
			this.showImages = showImages;
			this.invalidate();
		}
	}

	settle(): void {
		if (this.complete) return;
		this.complete = true;
		this.imageIds = this.preview.labels.map(() => allocateImageId());
		this.invalidate();
	}

	render(width: number): string[] {
		const version = this.preview.version();
		if (
			this.cachedLines &&
			this.cachedVersion === version &&
			this.cachedWidth === width
		)
			return this.cachedLines;
		if (this.preview.expired()) {
			this.cachedLines = [this.color("muted", "Preview expired.")];
			return this.cachedLines;
		}
		if (!this.showImages || !getCapabilities().images) {
			this.cachedLines = [
				this.color("warning", "Inline image previews are disabled."),
			];
			return this.cachedLines;
		}
		const frames = this.preview.current();
		if (!frames) {
			this.cachedLines = [this.color("accent", "Preparing previews…")];
			return this.cachedLines;
		}
		const cropped = this.preview.cropped();
		const maxHeightCells =
			frames.length === 1 ? 20 : frames.length === 2 ? 12 : 8;
		const lines: string[] = [];
		for (let index = 0; index < frames.length; index++) {
			if (index > 0) lines.push("");
			lines.push(
				this.color(
					"accent",
					this.preview.labels[index] ?? `Option ${index + 1}`,
				),
			);
			if (cropped[index])
				lines.push(
					this.color("warning", "Preview cropped · use browser review"),
				);
			const image = new Image(
				frames[index] ?? "",
				"image/png",
				{ fallbackColor: (text) => this.color("muted", text) },
				{
					imageId: this.imageIds[index],
					maxWidthCells: Math.min(80, Math.max(1, width - 2)),
					maxHeightCells,
				},
			);
			lines.push(...image.render(width));
		}
		this.cachedLines = lines;
		this.cachedVersion = version;
		this.cachedWidth = width;
		return lines;
	}

	invalidate(): void {
		this.cachedLines = undefined;
		this.cachedVersion = -1;
		this.cachedWidth = -1;
	}

	dispose(): void {
		this.unsubscribe();
	}
}

export function renderMockupPreviewPage(request: MockupRequest): string {
	const cards = request.options
		.map(
			(option) => `<article class="preview-card">
<div class="fragment-content"></div>
<script type="application/json" class="fragment-html">${jsonScript(option.html)}</script>
</article>`,
		)
		.join("\n");
	return `<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<style>
:root { color-scheme:dark light; --text:#ededed; --muted:#9a9a9a; --line:#ededed; font:14px/1.45 ui-monospace,SFMono-Regular,Menlo,Consolas,monospace; }
@media (prefers-color-scheme:light) { :root { --text:#151515; --muted:#666; --line:#151515; } }
* { box-sizing:border-box; }
html,body { margin:0; width:720px; background:transparent; color:var(--text); }
body { padding:20px; display:grid; gap:20px; }
.preview-card { width:680px; border:2px solid var(--line); background:transparent; overflow:hidden; }
.fragment-content { display:block; width:676px; min-height:300px; padding:16px; overflow:visible; }
</style>
</head>
<body>${cards}
<script>
(() => {
  document.querySelectorAll('.preview-card').forEach(card => {
    const content = card.querySelector('.fragment-content');
    const data = card.querySelector('.fragment-html');
    const root = content.attachShadow({ mode:'open' });
    root.innerHTML = JSON.parse(data.textContent || '""');
    root.querySelectorAll('script').forEach(old => {
      const script = document.createElement('script');
      [...old.attributes].forEach(attr => script.setAttribute(attr.name, attr.value));
      script.textContent = old.textContent;
      Object.defineProperty(script, 'mockupRoot', { value:root });
      const pending = Boolean(script.src) || script.type === 'module';
      if (pending) {
        script.addEventListener('load', () => script.remove(), { once:true });
        script.addEventListener('error', () => script.remove(), { once:true });
      }
      document.head.append(script);
      if (!pending) script.remove();
    });
    data.remove();
  });
  requestAnimationFrame(() => requestAnimationFrame(() => {
    const contents = [...document.querySelectorAll('.fragment-content')];
    const maximumContentHeight = ${MAX_PREVIEW_HEIGHT_PX - PREVIEW_BORDER_PX};
    const commonHeight = Math.min(
      maximumContentHeight,
      Math.max(300, ...contents.map(content => Math.ceil(content.getBoundingClientRect().height))),
    );
    contents.forEach(content => {
      content.style.height = commonHeight + 'px';
      content.style.overflow = 'hidden';
      const cropped = content.scrollHeight > content.clientHeight || content.scrollWidth > content.clientWidth;
      content.parentElement.toggleAttribute('data-preview-cropped', cropped);
    });
    document.documentElement.dataset.mockupPreviewReady = 'true';
  }));
})();
</script>
</body>
</html>`;
}

export function previewDeviceScaleFactor(
	optionCount: number,
	terminalColumns = 82,
	cells = getCellDimensions(),
): number {
	const maxHeightCells = optionCount === 1 ? 20 : optionCount === 2 ? 12 : 8;
	const widthScale =
		(Math.min(80, Math.max(1, terminalColumns - 2)) * cells.widthPx) /
		PREVIEW_WIDTH_PX;
	const heightScale = (maxHeightCells * cells.heightPx) / PREVIEW_HEIGHT_PX;
	return Math.min(2, Math.max(0.5, Math.min(widthScale, heightScale)));
}

export function thinReplayFrames(
	frames: readonly PreviewFrameSet[],
	maxFrames = MAX_REPLAY_FRAME_SETS,
	maxBytes = MAX_REPLAY_BYTES,
): PreviewFrameSet[] {
	let retained = [...frames];
	let bytes = retained.reduce((total, frame) => total + frameBytes(frame), 0);
	while (
		(retained.length > maxFrames || bytes > maxBytes) &&
		retained.length > 1
	) {
		if (retained.length === 2) {
			retained = [retained[1]];
		} else {
			const last = retained.at(-1);
			retained = retained.filter((_, index) => index % 2 === 0);
			if (last && retained.at(-1) !== last) retained.push(last);
		}
		bytes = retained.reduce((total, frame) => total + frameBytes(frame), 0);
	}
	return bytes <= maxBytes ? retained : [];
}

function frameBytes(frame: PreviewFrameSet): number {
	return frame.frames.reduce((total, image) => total + image.length, 0);
}

function sameFrames(
	left: readonly string[] | undefined,
	right: readonly string[],
): boolean {
	return (
		left?.length === right.length &&
		left.every((frame, index) => frame === right[index])
	);
}

function sameFlags(
	left: readonly boolean[] | undefined,
	right: readonly boolean[] | undefined,
): boolean {
	return (
		left?.length === right?.length &&
		(left?.every((flag, index) => flag === right?.[index]) ?? true)
	);
}

function jsonScript(value: string): string {
	return JSON.stringify(value).replaceAll("<", "\\u003c");
}