backend/omr-service/src/services/predict-pages.ts

/**
 * Thin wrapper around predictPages from starry-omr.
 *
 * Handles webp→png conversion, saves images to uploads dir,
 * and reports progress via the task service.
 */
import { createHash } from 'crypto';
import { existsSync } from 'fs';
import { writeFile } from 'fs/promises';
import { join } from 'path';
import sharp from 'sharp';
import type { LayoutResult, OMRPage, ProgressState } from 'starry-omr';
import { predictPages, PyClients } from 'starry-omr';
import { config } from '../config.js';
import * as scoreService from './score.service.js';
import * as taskService from './task.service.js';

export type { LayoutResult };

// skia-canvas loadImage(Buffer) only supports PNG/JPEG/GIF/BMP, NOT webp.
// Use sharp to convert unsupported formats to PNG.
async function ensurePng(buf: Buffer): Promise<Buffer> {
	// Check magic bytes: PNG starts with 89 50 4E 47, JPEG with FF D8
	if (buf[0] === 0x89 && buf[1] === 0x50) return buf; // PNG
	if (buf[0] === 0xff && buf[1] === 0xd8) return buf; // JPEG
	// Everything else (webp, etc.) — convert to PNG
	return sharp(buf).png().toBuffer();
}

const UPLOADS_DIR = process.env.UPLOADS_DIR || '/tmp/starry-uploads';

export interface PredictPagesOptions {
	taskId: string;
	scoreId: string;
	images: Array<{
		data: Buffer;
		layout?: LayoutResult;
		enableGauge?: boolean;
	}>;
	outputWidth?: number;
	processes?: string[];
}

// ─── Image persistence helpers ────────────────────────────────────────────

function md5Short(buf: Buffer): string {
	return createHash('md5').update(buf).digest('hex').slice(0, 12);
}

async function saveImageToUploads(buf: Buffer, prefix: string, ext = 'png'): Promise<string> {
	const hash = md5Short(buf);
	const filename = `${prefix}_${hash}.${ext}`;
	const filepath = join(UPLOADS_DIR, filename);
	if (!existsSync(filepath)) {
		await writeFile(filepath, buf);
	}
	return `/uploads/${filename}`;
}

// ─── Progress mapping ─────────────────────────────────────────────────────

function progressToPercent(state: ProgressState): number {
	// Map multi-stage progress to a 10–90% range
	const stages: (keyof ProgressState)[] = ['layout', 'gauge', 'semantic', 'mask', 'text', 'brackets'];
	let totalWeight = 0;
	let doneWeight = 0;

	for (const stage of stages) {
		const s = state[stage];
		if (s && s.total) {
			totalWeight += s.total;
			doneWeight += s.finished || 0;
		}
	}

	if (totalWeight === 0) return 10;
	return 10 + Math.floor((doneWeight / totalWeight) * 80);
}

// ─── Main pipeline ─────────────────────────────────────────────────────────

export async function predictPagesService(options: PredictPagesOptions): Promise<void> {
	const { taskId, scoreId, images, outputWidth = 1200, processes = ['semantic', 'mask', 'brackets', 'text'] } = options;

	try {
		await taskService.updateTaskStatus(taskId, 'running', 0, 'initializing');

		// Create PyClients from config addresses
		// PyClients keys must match PredictorInterface: textLoc/textOcr (not loc/ocr)
		const pyClients = new PyClients({
			layout: config.predictors.layout,
			semantic: config.predictors.semantic,
			mask: config.predictors.mask,
			gauge: config.predictors.gauge,
			gaugeRenderer: config.predictors.gaugeRenderer,
			textLoc: config.predictors.loc,
			textOcr: config.predictors.ocr,
			brackets: config.predictors.brackets,
		});

		// Build OMRPage[] from input buffers (with ensurePng for webp support)
		await taskService.updateTaskStatus(taskId, 'running', 5, 'loading images');
		const omrPages: OMRPage[] = await Promise.all(
			images.map(async (img) => ({
				url: await ensurePng(img.data),
				layout: img.layout,
				enableGauge: img.enableGauge,
			}))
		);

		console.log(`[predict-pages] Processing ${omrPages.length} pages...`);

		// Custom onReplaceImage: save Buffers to uploads dir, return /uploads/ URLs
		const prefix = `pp_${scoreId.slice(0, 8)}`;
		const onReplaceImage = async (src: string | Buffer): Promise<string> => {
			if (src instanceof Buffer) {
				return saveImageToUploads(src, prefix);
			}
			return src as string;
		};

		// Call the original predictPages — all coordinate logic lives there
		const { score } = await predictPages(pyClients, omrPages, {
			outputWidth,
			processes: processes as (keyof ProgressState)[],
			onReplaceImage,
			onProgress: (state) => {
				const pct = progressToPercent(state);
				taskService.updateTaskStatus(taskId, 'running', pct, 'predicting').catch(() => {});
			},
		});

		// Finalize
		await taskService.updateTaskStatus(taskId, 'running', 95, 'saving score');

		console.log(`[predict-pages] Score: ${score.systems.length} systems, staffLayout=${score.staffLayoutCode}`);

		// Save to DB
		await scoreService.updateScore(scoreId, { data: (score as any).toJSON() });

		// Complete task
		await taskService.setTaskResult(taskId, {
			systems: score.systems.length,
			staffLayout: score.staffLayoutCode,
			pageSize: score.pageSize,
			unitSize: score.unitSize,
		});

		console.log(`[predict-pages] Done! Score ${scoreId} saved.`);
	} catch (error) {
		const errorMessage = error instanceof Error ? error.message : String(error);
		console.error(`[predict-pages] Failed:`, error);
		await taskService.setTaskError(taskId, errorMessage);
		throw error;
	}
}