music-video-gen/flow-state/src/export/Exporter.js

import { Muxer, ArrayBufferTarget } from 'mp4-muxer';

/**
 * Offline export.
 *
 * Drives the same Show object the preview does, with the timeline in fixed-step
 * mode — the frame index is counted, never derived from a clock. That is the
 * whole basis of preview/export parity: the exporter has no render path of its
 * own, so there is nothing for it to diverge from.
 *
 * Video goes through WebCodecs VideoEncoder (hardware accelerated where
 * available) into an mp4. Audio is re-encoded from the already-decoded PCM —
 * AAC where available, Opus otherwise — so the file carries the same samples
 * the analysis ran on.
 */

export const PRESETS = {
    '720p':  { width: 1280, height: 720,  bitrate: 8_000_000 },
    '1080p': { width: 1920, height: 1080, bitrate: 16_000_000 },
    '1440p': { width: 2560, height: 1440, bitrate: 28_000_000 },
    '4K':    { width: 3840, height: 2160, bitrate: 45_000_000 },
};

export function isSupported() {
    return typeof VideoEncoder !== 'undefined' && typeof VideoFrame !== 'undefined';
}

/** Probe for a codec configuration the browser will actually accept. */
async function pickVideoConfig(width, height, bitrate, fps) {
    const candidates = [
        'avc1.640034', 'avc1.640033', 'avc1.4d0034', 'avc1.42003e',
    ];
    for (const codec of candidates) {
        const config = {
            codec, width, height, bitrate, framerate: fps,
            avc: { format: 'avc' },
        };
        try {
            const support = await VideoEncoder.isConfigSupported(config);
            if (support.supported) return config;
        } catch { /* try the next one */ }
    }
    return null;
}

/**
 * Pick an audio codec the browser can actually encode.
 *
 * AAC is the obvious choice for mp4 but is absent from Chromium builds without
 * proprietary codecs — which still ship H.264 encoding, so video succeeds and
 * only audio fails. Opus in mp4 is well supported by players and by every
 * platform worth uploading to, so it is the fallback rather than an error.
 */
async function pickAudioConfig(sampleRate, numberOfChannels) {
    const candidates = [
        { codec: 'mp4a.40.2', muxerCodec: 'aac', bitrate: 192_000 },
        { codec: 'opus', muxerCodec: 'opus', bitrate: 160_000 },
    ];
    for (const candidate of candidates) {
        try {
            const support = await AudioEncoder.isConfigSupported({
                codec: candidate.codec, sampleRate, numberOfChannels, bitrate: candidate.bitrate,
            });
            if (support.supported) return candidate;
        } catch { /* try the next one */ }
    }
    return null;
}

/**
 * Opus only produces a native 48 kHz stream (the input rate is resampled
 * inside the encoder), so a 44.1 kHz source would end up in a file whose
 * container metadata claims 44100 while the bitstream is 48000. Every stream
 * rate — encoder config, chunk timestamps and the muxer timescale — must
 * agree, so resample to the codec's native rate instead of trusting the
 * encoder to paper over the mismatch.
 */
export const OPUS_NATIVE_RATE = 48000;

/** Linear-interpolation resampler for interleaved PCM. */
function resampleInterleaved(input, fromRate, toRate, channels) {
    if (fromRate === toRate) return input;
    const outFrames = Math.floor((input.length / channels) * (toRate / fromRate));
    const out = new Float32Array(outFrames * channels);
    const ratio = fromRate / toRate;
    const last = input.length / channels - 1;
    for (let s = 0; s < outFrames; s++) {
        const pos = s * ratio;
        const i0 = Math.floor(pos);
        const i1 = Math.min(i0 + 1, last);
        const frac = pos - i0;
        for (let c = 0; c < channels; c++) {
            out[s * channels + c] =
                input[i0 * channels + c] * (1 - frac) + input[i1 * channels + c] * frac;
        }
    }
    return out;
}

/**
 * Backpressure: VideoEncoder.encode() silently drops frames once its internal
 * queue is full (encodeQueueSize is capped at an implementation-defined
 * limit). Polling that number is unreliable — the only signal that a frame
 * actually left the queue is the 'dequeue' event, so wait on it before every
 * encode. Keeping the queue small (rather than draining it fully) leaves the
 * hardware encoder a pipeline to chew on while guaranteeing nothing is dropped.
 */
const MAX_ENCODER_QUEUE = 8;

function waitForQueue(encoder, errors) {
    if (encoder.encodeQueueSize < MAX_ENCODER_QUEUE) return Promise.resolve();
    return new Promise((resolve, reject) => {
        const cleanup = () => {
            encoder.removeEventListener('dequeue', onDequeue);
            encoder.removeEventListener('error', onError);
        };
        const onError = () => { cleanup(); reject(errors[0] || new Error('encoder error')); };
        const onDequeue = () => {
            if (errors.length) { cleanup(); reject(errors[0]); return; }
            if (encoder.encodeQueueSize < MAX_ENCODER_QUEUE) { cleanup(); resolve(); }
        };
        encoder.addEventListener('dequeue', onDequeue);
        encoder.addEventListener('error', onError);
    });
}

export class Exporter {
    constructor(show) {
        this.show = show;
        this.cancelled = false;
    }

    cancel() { this.cancelled = true; }

    /**
     * @param {object} options
     * @param {string} options.preset key of PRESETS
     * @param {[number, number]} [options.frameRange] inclusive-exclusive frame range
     * @param {(progress: {frame, total, fraction, stage}) => void} [options.onProgress]
     * @returns {Promise<Blob>}
     */
    async export({ preset = '1080p', frameRange = null, onProgress = null } = {}) {
        if (!isSupported()) {
            throw new Error('WebCodecs VideoEncoder is unavailable in this browser');
        }
        const show = this.show;
        if (!show.ready) throw new Error('no track loaded');

        const { width, height, bitrate } = PRESETS[preset] || PRESETS['1080p'];
        const fps = show.fps;

        const [startFrame, endFrame] = frameRange || [0, show.frameCount];
        const total = Math.max(1, endFrame - startFrame);

        const videoConfig = await pickVideoConfig(width, height, bitrate, fps);
        if (!videoConfig) throw new Error('no supported H.264 configuration found');

        const channels = show.audioBuffer ? Math.min(2, show.audioBuffer.numberOfChannels) : 0;
        const audioConfig = show.audioBuffer && typeof AudioEncoder !== 'undefined'
            ? await pickAudioConfig(show.audioBuffer.sampleRate, channels)
            : null;
        const hasAudio = !!audioConfig;
        this.warnings = [];
        if (show.audioBuffer && !hasAudio) {
            this.warnings.push('no supported audio encoder — exporting video only');
        }

        // Every stream rate must agree so container metadata, chunk timestamps
        // and the bitstream describe the same timeline. Opus is resampled to
        // its native 48 kHz; AAC keeps the decoded PCM's rate.
        const sourceRate = show.audioBuffer ? show.audioBuffer.sampleRate : null;
        const audioRate = hasAudio && audioConfig.muxerCodec === 'opus'
            ? OPUS_NATIVE_RATE
            : sourceRate;

        const muxer = new Muxer({
            target: new ArrayBufferTarget(),
            video: { codec: 'avc', width, height, frameRate: fps },
            ...(hasAudio ? {
                audio: {
                    codec: audioConfig.muxerCodec,
                    sampleRate: audioRate,
                    numberOfChannels: channels,
                },
            } : {}),
            fastStart: 'in-memory',
        });

        const errors = [];
        const videoEncoder = new VideoEncoder({
            output: (chunk, meta) => muxer.addVideoChunk(chunk, meta),
            error: (e) => errors.push(e),
        });
        videoEncoder.configure(videoConfig);

        // Render at export resolution. The preview's own size is restored after.
        const previousWidth = show.engine.width;
        const previousHeight = show.engine.height;
        show.setSize(width, height);

        try {
            // Compile every shader and discard a warm frame first. Programs link
            // asynchronously, and an export renders each frame exactly once — there
            // is no second pass to fix frame 0 with.
            onProgress && onProgress({ frame: 0, total, fraction: 0, stage: 'compiling shaders' });
            show.prime(startFrame);

            // Warm-up so the first exported frame has the same feedback state it
            // would have had in sequential playback from the range start.
            if (startFrame > 0) {
                onProgress && onProgress({ frame: 0, total, fraction: 0, stage: 'warming up' });
                show.warmUp(startFrame, show.warmupFrames());
            } else {
                show.engine.compositor.reset();
            }

            for (let i = 0; i < total; i++) {
                if (this.cancelled) throw new Error('export cancelled');

                const frameIndex = startFrame + i;
                const target = show.renderFrame(frameIndex);
                show.present(target);   // encode from the canvas, which now holds this frame

                // True backpressure before every encode — VideoEncoder drops
                // frames silently once its queue saturates, which would turn a
                // smooth render into the stutter that first shipped here.
                await waitForQueue(videoEncoder, errors);
                if (errors.length) throw errors[0];

                const timestamp = Math.round(((frameIndex - startFrame) * 1e6) / fps);
                const videoFrame = new VideoFrame(show.engine.renderer.canvas, {
                    timestamp,
                    duration: Math.round(1e6 / fps),
                });
                // Keyframe every two seconds: seekable output without bloating size.
                videoEncoder.encode(videoFrame, { keyFrame: i % (fps * 2) === 0 });
                videoFrame.close();

                // Yield periodically so the progress UI paints.
                if (i % 10 === 0) {
                    onProgress && onProgress({
                        frame: i, total, fraction: i / total, stage: 'rendering',
                    });
                    await new Promise((r) => setTimeout(r, 0));
                }
                if (errors.length) throw errors[0];
            }

            onProgress && onProgress({ frame: total, total, fraction: 1, stage: 'finishing video' });
            await videoEncoder.flush();

            if (hasAudio) {
                onProgress && onProgress({ frame: total, total, fraction: 1, stage: 'encoding audio' });
                try {
                    await this._encodeAudio(muxer, audioConfig, startFrame, endFrame, fps, audioRate);
                } catch (err) {
                    // A finished silent video beats losing a long render outright.
                    // The muxer tolerates an audio track that received no chunks.
                    this.warnings.push(`audio encoding failed (${err.message}) — video only`);
                }
            }

            muxer.finalize();
            return new Blob([muxer.target.buffer], { type: 'video/mp4' });
        } finally {
            try { videoEncoder.close(); } catch { /* already closed */ }
            show.setSize(previousWidth, previousHeight);
        }
    }

    /** Encode the exported time range of the decoded PCM and mux it. */
    async _encodeAudio(muxer, audioConfig, startFrame, endFrame, fps, audioRate) {
        const buffer = this.show.audioBuffer;
        const sourceRate = buffer.sampleRate;
        const channels = Math.min(2, buffer.numberOfChannels);

        const startSample = Math.floor((startFrame / fps) * sourceRate);
        const endSample = Math.min(buffer.length, Math.ceil((endFrame / fps) * sourceRate));
        const length = Math.max(0, endSample - startSample);
        if (!length) return;

        const errors = [];
        const encoder = new AudioEncoder({
            output: (chunk, meta) => muxer.addAudioChunk(chunk, meta),
            error: (e) => errors.push(e),
        });
        encoder.configure({
            codec: audioConfig.codec,
            sampleRate: audioRate,
            numberOfChannels: channels,
            bitrate: audioConfig.bitrate,
        });

        const chunkFrames = 1024;
        const sources = [];
        for (let c = 0; c < channels; c++) sources.push(buffer.getChannelData(c));

        // Slice the range, resampling to the codec's native rate when needed so
        // chunk timestamps and the muxer timescale describe the same timeline.
        const interleaved = new Float32Array(length * channels);
        for (let i = 0; i < length; i++) {
            for (let c = 0; c < channels; c++) {
                interleaved[i * channels + c] = sources[c][startSample + i];
            }
        }
        const pcm = audioRate === sourceRate
            ? interleaved
            : resampleInterleaved(interleaved, sourceRate, audioRate, channels);

        const outFrames = pcm.length / channels;
        const buf = new Float32Array(chunkFrames * channels);

        for (let offset = 0; offset < outFrames; offset += chunkFrames) {
            const count = Math.min(chunkFrames, outFrames - offset);
            buf.set(pcm.subarray(offset * channels, (offset + count) * channels), 0);
            const data = new AudioData({
                format: 'f32',
                sampleRate: audioRate,
                numberOfFrames: count,
                numberOfChannels: channels,
                timestamp: Math.round((offset / audioRate) * 1e6),
                data: buf.slice(0, count * channels),
            });
            encoder.encode(data);
            data.close();
            if (errors.length) throw errors[0];
            if (offset % (chunkFrames * 64) === 0) await new Promise((r) => setTimeout(r, 0));
        }
        await encoder.flush();
        encoder.close();
    }
}

/** Render a short range around a frame — the "test render" bridge before a full export. */
export async function exportSegment(show, centreFrame,
    { seconds = 20, preset = '1080p', onProgress, exporter = null } = {}) {
    const half = Math.round((seconds * show.fps) / 2);
    const start = Math.max(0, centreFrame - half);
    const end = Math.min(show.frameCount, centreFrame + half);
    // Accept a caller-supplied Exporter so the caller can read `warnings` and cancel.
    return (exporter || new Exporter(show)).export({ preset, frameRange: [start, end], onProgress });
}

export function downloadBlob(blob, filename) {
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = filename;
    document.body.appendChild(a);
    a.click();
    a.remove();
    setTimeout(() => URL.revokeObjectURL(url), 10000);
}