// Background removal for camera tracks with Robust Video Matting (Lin et al., 2021), run on the CPU
// through onnxruntime. For each take, camera.webm is decoded to 30 fps RGB frames, the recurrent
// model predicts a foreground color and an alpha matte per frame, and the result is written as
// cutout.webm (VP9 with an alpha channel), which Chrome plays with transparency and ffmpeg can
// composite at export. Takes are processed one at a time, since each job already uses most cores.

import { spawn } from 'node:child_process';
import { once } from 'node:events';
import fs from 'node:fs/promises';
import os from 'node:os';
import path from 'node:path';
import type { Readable } from 'node:stream';
import ort from 'onnxruntime-node';
import type { Take, Video } from '../shared/types.ts';
import { FPS } from '../shared/layout.ts';
import { probe } from './media.ts';
import { HttpError, exists, getVideo, takeDir, updateVideo } from './store.ts';

const MODEL_URL = 'https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp32.onnx';
const MODEL_PATH = path.join(os.homedir(), '.cache', 'takes', 'models', 'rvm_mobilenetv3_fp32.onnx');

interface Job {
  videoId: string;
  takeId: string;
  progress: number;
  running: boolean;
}

const jobs = new Map<string, Job>(); // key: videoId/takeId, in queue order
let pumping = false;
let session: Promise<ort.InferenceSession> | null = null;

const key = (videoId: string, takeId: string) => `${videoId}/${takeId}`;

export const cutoutFile = (videoId: string, takeId: string) => path.join(takeDir(videoId, takeId), 'cutout.webm');

// Merge the live queue state into a video before sending it to the client.
export function withMatteStatus(v: Video): Video {
  const takes = v.takes.map((t): Take => {
    const j = jobs.get(key(v.id, t.id));
    return j ? { ...t, matte: { status: j.running ? 'running' : 'queued', progress: j.progress, updatedAt: new Date().toISOString() } } : t;
  });
  return { ...v, takes };
}

// Queue every ready camera take of the video that has not been processed yet (or failed, if retry).
export async function requestMattes(videoId: string, opts: { retry?: boolean; takeId?: string } = {}) {
  const v = await getVideo(videoId);
  for (const t of v.takes) {
    if (opts.takeId && t.id !== opts.takeId) continue;
    if (t.status !== 'ready' || !t.tracks.camera || jobs.has(key(videoId, t.id))) continue;
    if (t.matte?.status === 'done' && (await exists(cutoutFile(videoId, t.id)))) continue;
    if (t.matte?.status === 'failed' && !opts.retry) continue;
    jobs.set(key(videoId, t.id), { videoId, takeId: t.id, progress: 0, running: false });
  }
  pump();
}

async function pump() {
  if (pumping) return;
  pumping = true;
  try {
    for (;;) {
      const job = jobs.values().next().value as Job | undefined;
      if (!job) break;
      job.running = true;
      const setMatte = (matte: Take['matte']) =>
        updateVideo(job.videoId, (v) => {
          const t = v.takes.find((t) => t.id === job.takeId);
          if (t) t.matte = matte;
        }).catch(() => {}); // the video or take may have been deleted meanwhile
      try {
        const t0 = Date.now();
        const frames = await processTake(job);
        console.log(`matte: ${job.videoId}/${job.takeId}: ${frames} frames in ${((Date.now() - t0) / 1000).toFixed(1)} s`);
        await setMatte({ status: 'done', updatedAt: new Date().toISOString() });
      } catch (e) {
        console.warn(`matte failed for ${job.videoId}/${job.takeId}:`, (e as Error).message);
        await setMatte({ status: 'failed', error: (e as Error).message, updatedAt: new Date().toISOString() });
      } finally {
        jobs.delete(key(job.videoId, job.takeId));
      }
    }
  } finally {
    pumping = false;
  }
}

async function ensureModel() {
  if (await exists(MODEL_PATH)) return;
  console.log(`Downloading background removal model to ${MODEL_PATH}`);
  const res = await fetch(MODEL_URL);
  if (!res.ok) throw new Error(`Could not download the model (${res.status}) from ${MODEL_URL}`);
  await fs.mkdir(path.dirname(MODEL_PATH), { recursive: true });
  const tmp = `${MODEL_PATH}.${process.pid}.tmp`;
  await fs.writeFile(tmp, Buffer.from(await res.arrayBuffer()));
  await fs.rename(tmp, MODEL_PATH);
}

function getSession() {
  if (!session) {
    session = ensureModel().then(() =>
      ort.InferenceSession.create(MODEL_PATH, {
        executionProviders: ['cpu'],
        graphOptimizationLevel: 'all',
        intraOpNumThreads: Math.max(1, os.cpus().length - 2),
      }),
    );
    session.catch(() => (session = null)); // allow a retry after e.g. a failed download
  }
  return session;
}

// Yield fixed-size frames from a raw video stream. The same buffer is reused, so each frame must be
// consumed before asking for the next.
async function* readFrames(stream: Readable, size: number) {
  const frame = Buffer.alloc(size);
  let off = 0;
  for await (const chunk of stream as AsyncIterable<Buffer>) {
    let p = 0;
    while (p < chunk.length) {
      const n = Math.min(size - off, chunk.length - p);
      chunk.copy(frame, off, p, p + n);
      off += n;
      p += n;
      if (off === size) {
        yield frame;
        off = 0;
      }
    }
  }
}

const even = (n: number) => Math.max(2, Math.round(n / 2) * 2);

async function processTake(job: Job): Promise<number> {
  const sess = await getSession();
  const src = path.join(takeDir(job.videoId, job.takeId), 'camera.webm');
  const info = await probe(src);
  if (!info.width || !info.height) throw new HttpError(400, 'Camera track has no video');
  // Process at the camera's resolution, capped at 1080 lines.
  const scale = Math.min(1, 1080 / info.height);
  const W = even(info.width * scale);
  const H = even(info.height * scale);
  // Downsample ratio for the model's coarse pass, following the RVM guidance (0.25 at 1080p, 0.375 at 720p).
  const ds = Math.min(1, 480 / Math.max(W, H));
  const total = Math.max(1, Math.round(info.duration * FPS));

  const out = cutoutFile(job.videoId, job.takeId);
  const tmp = `${out}.partial.webm`;
  const decoder = spawn('ffmpeg', ['-v', 'error', '-i', src, '-an', '-vf', `fps=${FPS}:start_time=0,scale=${W}:${H}`, '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-'], { stdio: ['ignore', 'pipe', 'pipe'] });
  const encoder = spawn(
    'ffmpeg',
    ['-y', '-v', 'error', '-f', 'rawvideo', '-pix_fmt', 'rgba', '-s', `${W}x${H}`, '-r', String(FPS), '-i', '-',
      '-c:v', 'libvpx-vp9', '-pix_fmt', 'yuva420p', '-deadline', 'realtime', '-cpu-used', '8', '-row-mt', '1', '-b:v', '3M', '-auto-alt-ref', '0', tmp],
    { stdio: ['pipe', 'ignore', 'pipe'] },
  );
  let stderr = '';
  decoder.stderr.on('data', (d) => (stderr += d));
  encoder.stderr.on('data', (d) => (stderr += d));
  const decoderDone = once(decoder, 'close');
  const encoderDone = once(encoder, 'close');

  const plane = W * H;
  const input = new Float32Array(3 * plane);
  let rec: ort.Tensor[] = [0, 1, 2, 3].map(() => new ort.Tensor('float32', new Float32Array([0]), [1, 1, 1, 1]));
  const dsr = new ort.Tensor('float32', new Float32Array([ds]), [1]);
  let n = 0;
  try {
    for await (const frame of readFrames(decoder.stdout, plane * 3)) {
      for (let i = 0, j = 0; i < plane; i++, j += 3) {
        input[i] = frame[j] / 255;
        input[i + plane] = frame[j + 1] / 255;
        input[i + 2 * plane] = frame[j + 2] / 255;
      }
      const r = await sess.run({
        src: new ort.Tensor('float32', input, [1, 3, H, W]),
        r1i: rec[0],
        r2i: rec[1],
        r3i: rec[2],
        r4i: rec[3],
        downsample_ratio: dsr,
      });
      rec = [r.r1o, r.r2o, r.r3o, r.r4o];
      const fgr = r.fgr.data as Float32Array;
      const pha = r.pha.data as Float32Array;
      const rgba = new Uint8ClampedArray(plane * 4);
      for (let i = 0, j = 0; i < plane; i++, j += 4) {
        rgba[j] = fgr[i] * 255;
        rgba[j + 1] = fgr[i + plane] * 255;
        rgba[j + 2] = fgr[i + 2 * plane] * 255;
        rgba[j + 3] = pha[i] * 255;
      }
      if (!encoder.stdin.write(Buffer.from(rgba.buffer))) await once(encoder.stdin, 'drain');
      n++;
      job.progress = Math.min(0.99, n / total);
    }
    encoder.stdin.end();
    const [[dcode], [ecode]] = await Promise.all([decoderDone, encoderDone]);
    if (dcode !== 0 || ecode !== 0) throw new Error(`ffmpeg failed: ${stderr.trim().split('\n').slice(-4).join('\n')}`);
    if (n === 0) throw new Error('No camera frames were decoded');
    await fs.rename(tmp, out);
    return n;
  } catch (e) {
    decoder.kill();
    encoder.kill();
    await fs.rm(tmp, { force: true });
    throw e;
  }
}