backend/omr-service/src/routes/predict.ts

import { MultipartFile } from '@fastify/multipart';
import { FastifyInstance } from 'fastify';
import * as scoreService from '../services/score.service.js';
import * as taskService from '../services/task.service.js';
import { taskWorker } from '../worker.js';

interface PredictPageBody {
	score_id: string;
	page_index: number;
	image: MultipartFile;
}

interface PredictAllBody {
	score_id: string;
}

interface PredictPageParams {
	scoreId: string;
	pageIndex: string;
}

export default async function predictRoutes(fastify: FastifyInstance) {
	// Predict single page (async task)
	fastify.post<{ Params: PredictPageParams }>('/predict/page/:scoreId/:pageIndex', async (request, reply) => {
		const { scoreId, pageIndex } = request.params;
		const pageIdx = parseInt(pageIndex, 10);

		// Check score exists
		const score = await scoreService.getScore(scoreId);
		if (!score) {
			reply.code(404);
			return { error: 'Score not found' };
		}

		// Handle multipart file upload
		const data = await request.file();
		if (!data) {
			reply.code(400);
			return { error: 'Image file required' };
		}

		const imageBuffer = await data.toBuffer();

		// Create page record if not exists
		await scoreService.createPage(scoreId, pageIdx, null, null, null, {});

		// Create async task
		const task = await taskService.createTask({
			score_id: scoreId,
			type: 'predict_page',
		});

		// Queue the task for processing
		taskWorker.queueTask({
			taskId: task.id,
			scoreId,
			pageIndex: pageIdx,
			imageData: imageBuffer,
		});

		reply.code(202);
		return {
			task_id: task.id,
			status: 'pending',
			poll_url: `/api/tasks/${task.id}/poll`,
		};
	});

	// Predict all pages (async task)
	fastify.post('/predict/all/:scoreId', async (request, reply) => {
		const scoreId = (request.params as { scoreId: string }).scoreId;

		// Check score exists
		const score = await scoreService.getScore(scoreId);
		if (!score) {
			reply.code(404);
			return { error: 'Score not found' };
		}

		// Handle multipart with multiple files
		const parts = request.files();
		const imageBuffers: Buffer[] = [];

		for await (const part of parts) {
			imageBuffers.push(await part.toBuffer());
		}

		if (imageBuffers.length === 0) {
			reply.code(400);
			return { error: 'At least one image file required' };
		}

		// Create page records
		for (let i = 0; i < imageBuffers.length; i++) {
			await scoreService.createPage(scoreId, i, null, null, null, {});
		}

		// Create async task
		const task = await taskService.createTask({
			score_id: scoreId,
			type: 'predict_all',
		});

		// Queue the task for processing
		taskWorker.queueTask({
			taskId: task.id,
			scoreId,
			imageDataList: imageBuffers,
		});

		reply.code(202);
		return {
			task_id: task.id,
			status: 'pending',
			page_count: imageBuffers.length,
			poll_url: `/api/tasks/${task.id}/poll`,
		};
	});

	// Predict pages using the full predictPages pipeline (Canvas-based rotation correction + staff extraction)
	// Produces a complete starry.Score with correct coordinates — no manual merging needed.
	// Accepts multipart form: page images (files) + 'sources' JSON field + optional 'outputWidth' + 'processes'
	fastify.post('/predict/pages/:scoreId', async (request, reply) => {
		const scoreId = (request.params as { scoreId: string }).scoreId;

		const score = await scoreService.getScore(scoreId);
		if (!score) {
			reply.code(404);
			return { error: 'Score not found' };
		}

		const imageBuffers: Buffer[] = [];
		let sourcesList: any[] = [];
		let outputWidth = 1200;
		let processes = ['semantic', 'mask'];

		const parts = request.parts();
		for await (const part of parts) {
			if (part.type === 'file') {
				imageBuffers.push(await (part as any).toBuffer());
			} else if (part.type === 'field') {
				if (part.fieldname === 'sources') {
					try {
						sourcesList = JSON.parse(part.value as string);
					} catch {
						reply.code(400);
						return { error: 'Invalid sources JSON' };
					}
				} else if (part.fieldname === 'outputWidth') {
					outputWidth = parseInt(part.value as string, 10) || 1200;
				} else if (part.fieldname === 'processes') {
					try {
						processes = JSON.parse(part.value as string);
					} catch {
						// keep default
					}
				}
			}
		}

		if (imageBuffers.length === 0) {
			reply.code(400);
			return { error: 'At least one image file required' };
		}

		const task = await taskService.createTask({
			score_id: scoreId,
			type: 'predict_custom',
		});

		const images = imageBuffers.map((data, i) => ({
			data,
			layout: sourcesList[i]?.layout || undefined,
			enableGauge: sourcesList[i]?.enableGauge || false,
		}));

		taskWorker.queueTask({
			taskId: task.id,
			scoreId,
			predictPages: { images, outputWidth, processes },
		});

		reply.code(202);
		return {
			task_id: task.id,
			status: 'pending',
			page_count: imageBuffers.length,
			poll_url: `/api/tasks/${task.id}/poll`,
		};
	});

	// Predict with pre-computed layout data (skip layout detection, only run semantic + mask)
	// Accepts multipart form: page images (files) + layout JSON (field named 'layout')
	fastify.post('/predict/with-layout/:scoreId', async (request, reply) => {
		const scoreId = (request.params as { scoreId: string }).scoreId;

		const score = await scoreService.getScore(scoreId);
		if (!score) {
			reply.code(404);
			return { error: 'Score not found' };
		}

		// Parse multipart: files are page images, 'layout' field is JSON
		const imageBuffers: Buffer[] = [];
		let layoutList: any[] = [];

		const parts = request.parts();
		for await (const part of parts) {
			if (part.type === 'file') {
				imageBuffers.push(await (part as any).toBuffer());
			} else if (part.type === 'field' && part.fieldname === 'layout') {
				try {
					layoutList = JSON.parse(part.value as string);
				} catch {
					reply.code(400);
					return { error: 'Invalid layout JSON' };
				}
			}
		}

		if (imageBuffers.length === 0) {
			reply.code(400);
			return { error: 'At least one image file required' };
		}

		// Create page records
		for (let i = 0; i < imageBuffers.length; i++) {
			await scoreService.createPage(scoreId, i, null, null, null, {});
		}

		const task = await taskService.createTask({
			score_id: scoreId,
			type: 'predict_all',
		});

		taskWorker.queueTask({
			taskId: task.id,
			scoreId,
			imageDataList: imageBuffers,
			layoutList,
		});

		reply.code(202);
		return {
			task_id: task.id,
			status: 'pending',
			page_count: imageBuffers.length,
			poll_url: `/api/tasks/${task.id}/poll`,
		};
	});
}