backend/omr-service/src/services/measure.service.ts

import { existsSync } from 'fs';
import { join } from 'path';
import { starry } from 'starry-omr';
import { transaction } from '../db/client.js';
import { getPickers, isRegulationReady } from '../lib/regulation.js';
import * as issueMeasureService from './issueMeasure.service.js';
import { getScore } from './score.service.js';
import * as solutionCacheService from './solutionCache.service.js';

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

interface MeasureSummary {
	measureIndex: number;
	eventCount: number;
	voiceCount: number;
	regulated: boolean;
	tickTwist: number | null;
	timeSignature: { numerator: number; denominator: number } | null;
	regulationHash: string;
	qualityScore: number | null;
}

interface MeasureIssue {
	measureIndex: number;
	tickTwist: number | null;
	qualityScore: number;
	error: boolean;
	fine: boolean;
	eventCount: number;
}

async function recoverScore(scoreData: any) {
	const score = starry.recoverJSON(scoreData, starry);
	(score as any).spartito = undefined;
	(score as any).assemble();
	const spartito = (score as any).makeSpartito();
	spartito.measures.forEach((measure: any) => (score as any).assignBackgroundForMeasure(measure));

	// Load cached solutions for non-patched measures that have events but no regulation
	const unregulated = spartito.measures.filter((m: any) => m.events?.length && !m.regulated);
	if (unregulated.length > 0) {
		const hashes = unregulated.map((m: any) => m.regulationHash);
		const solutions = await solutionCacheService.batchGet(hashes);
		for (let i = 0; i < unregulated.length; i++) {
			if (solutions[i]) {
				unregulated[i].applySolution(solutions[i]);
			}
		}
	}

	return { score, spartito };
}

function measureToSummary(measure: any): MeasureSummary {
	const evaluation = starry.evaluateMeasure(measure);
	return {
		measureIndex: measure.measureIndex,
		eventCount: measure.events?.length ?? 0,
		voiceCount: measure.voices?.length ?? 0,
		regulated: !!measure.regulated,
		tickTwist: evaluation?.tickTwist ?? null,
		timeSignature: measure.timeSignature ? { numerator: measure.timeSignature.numerator, denominator: measure.timeSignature.denominator } : null,
		regulationHash: measure.regulationHash,
		qualityScore: evaluation?.qualityScore ?? null,
	};
}

export async function getMeasures(scoreId: string) {
	const score = await getScore(scoreId);
	if (!score) return null;

	const { spartito } = await recoverScore(score.data);
	const measures = spartito.measures.map(measureToSummary);
	return { measures };
}

export async function getMeasure(scoreId: string, index: number) {
	const score = await getScore(scoreId);
	if (!score) return null;

	const { spartito } = await recoverScore(score.data);
	if (index < 0 || index >= spartito.measures.length) return null;

	const measure = spartito.measures[index];
	const evaluation = starry.evaluateMeasure(measure);

	return {
		measure: measure.toJSON(),
		evaluation: evaluation ?? null,
	};
}

export async function updateMeasure(scoreId: string, index: number, patch: any) {
	return transaction(async (client) => {
		const { rows } = await client.query('SELECT * FROM scores WHERE id = $1 FOR UPDATE', [scoreId]);
		if (!rows[0]) return null;

		const scoreRow = rows[0];
		const data = scoreRow.data;
		const { score, spartito } = await recoverScore(data);
		if (index < 0 || index >= spartito.measures.length) return null;

		// Apply patch fields to the measure
		const measure = spartito.measures[index];
		if (patch.events) measure.events = starry.recoverJSON(patch.events, starry);
		if (patch.voices) measure.voices = patch.voices;
		if (patch.duration !== undefined) measure.duration = patch.duration;
		if (patch.contexts) measure.contexts = starry.recoverJSON(patch.contexts, starry);
		if (patch.marks) measure.marks = starry.recoverJSON(patch.marks, starry);

		// Build a PatchMeasure and update score.patches
		const patchMeasure = new starry.PatchMeasure({
			measureIndex: index,
			staffMask: measure.staffMask,
			basic: measure.basics?.[0],
			events: measure.events,
			contexts: measure.contexts,
			marks: measure.marks,
			voices: measure.voices,
		});

		const patches = (score as any).patches || [];
		const existingIdx = patches.findIndex((p: any) => p.measureIndex === index);
		if (existingIdx >= 0) {
			patches[existingIdx] = patchMeasure;
		} else {
			patches.push(patchMeasure);
		}
		(score as any).patches = patches;

		// Write back
		const updatedData = (score as any).toJSON();
		await client.query('UPDATE scores SET data = $2, updated_at = NOW() WHERE id = $1', [scoreId, JSON.stringify(updatedData)]);

		const evaluation = starry.evaluateMeasure(measure);
		return {
			measure: measure.toJSON(),
			evaluation: evaluation ?? null,
		};
	});
}

export async function regulateMeasure(scoreId: string, index: number) {
	if (!isRegulationReady()) {
		throw new Error('Regulation service not ready');
	}

	const score = await getScore(scoreId);
	if (!score) return null;

	const { spartito } = await recoverScore(score.data);
	if (index < 0 || index >= spartito.measures.length) return null;

	const measure = spartito.measures[index];
	if (!measure.events?.length) {
		return { measure: measure.toJSON(), evaluation: null, solution: null };
	}

	const pickers = getPickers()!;
	const picker = pickers.find((p: any) => p.n_seq > measure.events.length + 1);
	if (!picker) {
		throw new Error(`No picker available for measure with ${measure.events.length} events`);
	}

	// Estimate time signature first
	await starry.beadSolver.estimateMeasure(measure, picker);

	// Try simple regulation first
	const simpleMeasure = measure.deepCopy();
	simpleMeasure.staffGroups = measure.staffGroups;
	simpleMeasure.regulate({ policy: 'simple' });
	const simpleEval = starry.evaluateMeasure(simpleMeasure);

	// Run bead solver
	const solution = await starry.beadSolver.solveMeasure(measure, {
		picker,
		stopLoss: 0.08,
		quotaMax: 1000,
		quotaFactor: 20,
		ptFactor: 1.6,
	});

	const solvedMeasure = measure.deepCopy();
	solvedMeasure.staffGroups = measure.staffGroups;
	solvedMeasure.applySolution(solution);
	const solvedEval = starry.evaluateMeasure(solvedMeasure);

	// Use the better result
	let bestMeasure, bestEval;
	if (simpleEval?.perfect) {
		bestMeasure = simpleMeasure;
		bestEval = simpleEval;
	} else if (!simpleEval || (solvedEval && solvedEval.qualityScore >= simpleEval.qualityScore)) {
		bestMeasure = solvedMeasure;
		bestEval = solvedEval;
	} else {
		bestMeasure = simpleMeasure;
		bestEval = simpleEval;
	}

	return {
		measure: bestMeasure.toJSON(),
		evaluation: bestEval ?? null,
		solution: {
			events: solution.events,
			voices: solution.voices,
			duration: solution.duration,
		},
	};
}

export async function getIssueMeasures(scoreId: string) {
	const score = await getScore(scoreId);
	if (!score) return null;

	const { spartito } = await recoverScore(score.data);
	const issues: MeasureIssue[] = [];

	for (const measure of spartito.measures) {
		if (!measure.events?.length) continue;

		const evaluation = starry.evaluateMeasure(measure);
		if (!evaluation) {
			// Unregulated measures are issues
			issues.push({
				measureIndex: measure.measureIndex,
				tickTwist: null,
				qualityScore: 0,
				error: false,
				fine: false,
				eventCount: measure.events.length,
			});
			continue;
		}

		if (!evaluation.fine || evaluation.error) {
			issues.push({
				measureIndex: measure.measureIndex,
				tickTwist: evaluation.tickTwist,
				qualityScore: evaluation.qualityScore,
				error: evaluation.error,
				fine: evaluation.fine,
				eventCount: measure.events.length,
			});
		}
	}

	return { issues };
}

export async function updateMeasureStatus(scoreId: string, index: number, status: number) {
	return transaction(async (client) => {
		const { rows } = await client.query('SELECT * FROM scores WHERE id = $1 FOR UPDATE', [scoreId]);
		if (!rows[0]) return null;

		const scoreRow = rows[0];
		const data = scoreRow.data;

		// Initialize measureStatuses if missing
		if (!data.measureStatuses) data.measureStatuses = {};
		data.measureStatuses[index] = status;

		await client.query('UPDATE scores SET data = $2, updated_at = NOW() WHERE id = $1', [scoreId, JSON.stringify(data)]);

		return { measureIndex: index, status };
	});
}

export interface AnnotateBody {
	fixes?: {
		clearGrace?: number[];
		setDivision?: { [eventIndex: string]: { division: number; dots: number } };
		setRest?: { [eventIndex: string]: boolean };
	};
	regulate?: boolean;
	policy?: string;
	status?: number;
	annotator?: string;
}

export async function annotateMeasure(scoreId: string, index: number, body: AnnotateBody) {
	const saved = { issueMeasure: false, solutionCache: false, scorePatch: false };

	// Step 1-7: Apply fixes, regulate, evaluate, save score patches (in transaction)
	const txResult = await transaction(async (client) => {
		const { rows } = await client.query('SELECT * FROM scores WHERE id = $1 FOR UPDATE', [scoreId]);
		if (!rows[0]) return null;

		const scoreRow = rows[0];
		const data = scoreRow.data;
		const { score, spartito } = await recoverScore(data);

		if (index < 0 || index >= spartito.measures.length) return { error: 'invalid_index' as const };

		const measure = spartito.measures[index];

		// Apply fixes to events
		if (body.fixes) {
			const { clearGrace, setDivision, setRest } = body.fixes;

			if (clearGrace) {
				for (const idx of clearGrace) {
					if (measure.events[idx]) {
						measure.events[idx].grace = null;
					}
				}
			}

			if (setDivision) {
				for (const [idxStr, value] of Object.entries(setDivision)) {
					const idx = parseInt(idxStr, 10);
					if (measure.events[idx]) {
						measure.events[idx].division = value.division;
						measure.events[idx].dots = value.dots;
					}
				}
			}

			if (setRest) {
				for (const [idxStr, value] of Object.entries(setRest)) {
					const idx = parseInt(idxStr, 10);
					if (measure.events[idx]) {
						measure.events[idx].rest = value ? 'R' : null;
					}
				}
			}
		}

		// Re-regulate (unless explicitly disabled)
		const policy = body.policy || 'equations';
		if (body.regulate !== false) {
			// Clear existing regulation state
			measure.voices = undefined;
			for (const event of measure.events || []) {
				event.tick = undefined;
				event.tickGroup = undefined;
				event.timeWarp = undefined;
			}

			if (policy === 'advanced') {
				if (!isRegulationReady()) {
					throw new Error('Regulation service not ready');
				}
				const pickers = getPickers()!;
				const picker = pickers.find((p: any) => p.n_seq > measure.events.length + 1);
				if (!picker) {
					throw new Error(`No picker available for measure with ${measure.events.length} events`);
				}
				await starry.beadSolver.estimateMeasure(measure, picker);
				const solution = await starry.beadSolver.solveMeasure(measure, {
					picker,
					stopLoss: 0.08,
					quotaMax: 1000,
					quotaFactor: 20,
					ptFactor: 1.6,
				});
				measure.applySolution(solution);
			} else {
				measure.regulate({ policy });
			}
		}

		// Evaluate
		const evaluation = starry.evaluateMeasure(measure);

		// Save score patches (inside transaction)
		const patchMeasure = new starry.PatchMeasure({
			measureIndex: index,
			staffMask: measure.staffMask,
			basic: measure.basics?.[0],
			events: measure.events,
			contexts: measure.contexts,
			marks: measure.marks,
			voices: measure.voices,
		});

		const patches = (score as any).patches || [];
		const existingIdx = patches.findIndex((p: any) => p.measureIndex === index);
		if (existingIdx >= 0) {
			patches[existingIdx] = patchMeasure;
		} else {
			patches.push(patchMeasure);
		}
		(score as any).patches = patches;

		const updatedData = (score as any).toJSON();
		await client.query('UPDATE scores SET data = $2, updated_at = NOW() WHERE id = $1', [scoreId, JSON.stringify(updatedData)]);
		saved.scorePatch = true;

		return { measure, evaluation };
	});

	if (!txResult) return null;
	if ('error' in txResult) return txResult;

	const { measure, evaluation } = txResult;

	// Save issue_measures (outside transaction)
	try {
		await issueMeasureService.upsert(scoreId, index, measure.toJSON(), body.status ?? 0, body.annotator);
		saved.issueMeasure = true;
	} catch (err) {
		console.error('[annotate] failed to save issue_measure:', err);
	}

	// Save solution_cache (if regulated and has solution)
	if (body.regulate !== false && measure.regulationHash) {
		try {
			const solution = measure.asSolution();
			if (solution) {
				await solutionCacheService.set(measure.regulationHash, solution);
				saved.solutionCache = true;
			}
		} catch (err) {
			console.error('[annotate] failed to save solution_cache:', err);
		}
	}

	return {
		measure: measure.toJSON(),
		evaluation: evaluation ?? null,
		saved,
	};
}

export async function getMeasureBackground(scoreId: string, index: number) {
	const score = await getScore(scoreId);
	if (!score) return null;

	const { spartito } = await recoverScore(score.data);
	if (index < 0 || index >= spartito.measures.length) return null;

	const measure = spartito.measures[index];
	const bgImages = measure.backgroundImages || [];
	const position = measure.position || {};

	// Build per-staff image info with crop coordinates
	const staffImages: {
		staff: number;
		url: string;
		filePath: string;
		exists: boolean;
		imagePosition: { x: number; y: number; width: number; height: number };
		measureCrop: { left: number; right: number; top: number; bottom: number; width: number; height: number };
	}[] = [];

	for (const bg of bgImages) {
		if (bg.original) continue; // Skip full-page images
		const url = bg.url || '';
		const filename = url.replace(/^\/uploads\//, '');
		const filePath = join(UPLOADS_DIR, filename);
		const exists = existsSync(filePath);
		const imgPos = bg.position || { x: 0, y: 0, width: 1, height: 1 };

		// Determine which staff this image belongs to
		// Image y is the top edge; staff center is near imgPos.y + imgPos.height / 2
		const staffYs = position.staffYs || [];
		const imgCenterY = imgPos.y + imgPos.height / 2;
		let staffIndex = staffImages.length; // fallback
		for (let i = 0; i < staffYs.length; i++) {
			if (Math.abs(imgCenterY - staffYs[i]) < 2) {
				staffIndex = i;
				break;
			}
		}

		// Calculate pixel crop for the measure region within this staff image
		// Image coordinate space: imgPos.x to imgPos.x + imgPos.width (in units)
		// We need actual pixel dimensions to compute the crop
		// Scale factor will be applied by the route handler after reading image metadata
		staffImages.push({
			staff: staffIndex,
			url,
			filePath,
			exists,
			imagePosition: imgPos,
			measureCrop: {
				left: position.left - imgPos.x,
				right: position.right - imgPos.x,
				top: 0,
				bottom: imgPos.height,
				width: position.right - position.left,
				height: imgPos.height,
			},
		});
	}

	return {
		measureIndex: index,
		measureBounds: { left: position.left, right: position.right },
		staffYs: position.staffYs || [],
		staffImages,
	};
}