import {
  AudioBufferSource,
  BufferTarget,
  CanvasSource,
  Mp4OutputFormat,
  Output,
  Quality,
  WebMOutputFormat,
} from "mediabunny";
import { extractFreqFrames } from "./input.ts";

type VisualizerOptions = Record<string, string | number | HTMLImageElement>;

export type OutputFormat = "webm" | "mp4";

export async function output(
  canvas: HTMLCanvasElement,
  buffer: AudioBuffer,
  draw: (
    canvas: HTMLCanvasElement,
    freq: Uint8Array,
    options: VisualizerOptions,
    dt?: number,
  ) => void,
  fps: number,
  quality: Quality,
  options: VisualizerOptions = {},
  audioReactive = true,
  format: OutputFormat = "webm",
  onProgress?: (current: number, total: number) => void,
) {
  const freqFrames = audioReactive
    ? await extractFreqFrames(buffer, fps)
    : null;
  const frameCount = audioReactive
    ? freqFrames!.length
    : Math.ceil(buffer.duration * fps);
  const EMPTY_FREQ = new Uint8Array(0);

  // move rendering to offscreen canvas for export to avoid
  // browser painting
  const offscreen = document.createElement("canvas");
  offscreen.width = canvas.width;
  offscreen.height = canvas.height;

  const isMp4 = format === "mp4";

  const output = new Output({
    format: isMp4 ? new Mp4OutputFormat() : new WebMOutputFormat(),
    target: new BufferTarget(),
  });

  const sourceVideo = new CanvasSource(offscreen, {
    codec: isMp4 ? "avc" : "vp9",
    bitrate: quality,
  });

  const sourceAudio = new AudioBufferSource({
    codec: isMp4 ? "aac" : "opus",
    bitrate: quality,
  });

  output.addVideoTrack(sourceVideo);
  output.addAudioTrack(sourceAudio);

  await output.start();

  for (let i = 0; i < frameCount; i++) {
    const freq = freqFrames ? freqFrames[i] : EMPTY_FREQ;
    draw(offscreen, freq, options, 1 / fps);
    await sourceVideo.add(i / fps, 1 / fps);
    onProgress?.(i + 1, frameCount);
  }
  sourceVideo.close();

  await sourceAudio.add(buffer);
  sourceAudio.close();

  await output.finalize();

  const finalBuffer = output.target.buffer!;
  return new Blob([finalBuffer], {
    type: isMp4 ? "video/mp4" : "video/webm",
  });
}
