// Background removal for camera tracks with Robust Video Matting (Lin et al., 2021), run on the CPU
// through onnxruntime, or on a remote machine's GPU over ssh when a host is set in the settings. 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 { createReadStream, createWriteStream } from 'node:fs';
import fs from 'node:fs/promises';
import os from 'node:os';
import path from 'node:path';
import type { Readable } from 'node:stream';
import { pipeline } from 'node:stream/promises';
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, getConfig, 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 { matteHost, matteSetup } = await getConfig();
        const frames = await processTake(job, { host: matteHost, setup: matteSetup });
        console.log(`matte: ${job.videoId}/${job.takeId}: ${frames} frames in ${((Date.now() - t0) / 1000).toFixed(1)} s${matteHost ? ` on ${matteHost}` : ''}`);
        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);

// ffmpeg arguments shared by the local and the remote path, so that both produce the same cutout.webm.
const decoderArgs = (src: string, W: number, H: number) => ['-v', 'error', '-i', src, '-an', '-vf', `fps=${FPS}:start_time=0,scale=${W}:${H}`, '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-'];
const encoderArgs = (out: string, W: number, H: number) =>
  ['-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', out];

interface Plan {
  src: string;
  out: string;
  W: number;
  H: number;
  ds: number;
  total: number;
}

async function processTake(job: Job, remote: { host: string; setup: string }): Promise<number> {
  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 plan: Plan = { src, out: tmp, W, H, ds, total };
  try {
    const n = remote.host ? await runRemote(job, plan, remote.host, remote.setup) : await runLocal(job, plan);
    if (n === 0) throw new Error('No camera frames were decoded');
    await fs.rename(tmp, out);
    return n;
  } catch (e) {
    await fs.rm(tmp, { force: true });
    throw e;
  }
}

async function runLocal(job: Job, { src, out, W, H, ds, total }: Plan): Promise<number> {
  const sess = await getSession();
  const decoder = spawn('ffmpeg', decoderArgs(src, W, H), { stdio: ['ignore', 'pipe', 'pipe'] });
  const encoder = spawn('ffmpeg', encoderArgs(out, W, H), { 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')}`);
    return n;
  } catch (e) {
    decoder.kill();
    encoder.kill();
    throw e;
  }
}

// ---- remote processing over ssh ----
//
// A single ssh connection carries the whole job: camera.webm goes to the remote machine on stdin,
// the remote shell script below runs matte-worker.py on it in a temporary directory, and the
// finished cutout.webm comes back on stdout. On first use the script creates a Python environment
// with onnxruntime-gpu in ~/.cache/takes/venv on the remote machine; pip installs the CUDA and cuDNN
// libraries as well, so the remote machine needs only an NVIDIA driver, Python 3.10 or later, and
// ffmpeg (possibly made available by the setup command, e.g. "module load python ffmpeg").

const WORKER = fs.readFile(new URL('./matte-worker.py', import.meta.url), 'utf8');

const shq = (s: string) => `'${s.replaceAll("'", `'\\''`)}'`;

function remoteScript(setup: string, worker: string) {
  return `{ ${setup || ':'}
} </dev/null
set -e
base=~/.cache/takes
venv=$base/venv
if [ ! -f "$venv/.ready" ]; then
  python3 -c 'import sys; sys.exit(sys.version_info < (3, 10))' </dev/null || {
    echo "Python 3.10 or later is needed on $(hostname), found: $(python3 --version 2>&1)" >&2; exit 1; }
  echo "Installing onnxruntime-gpu in $venv on $(hostname) (first use only; about 2.5 GB)" >&2
  rm -rf "$venv"
  mkdir -p "$base"
  python3 -m venv "$venv" </dev/null
  "$venv/bin/pip" install --quiet --disable-pip-version-check 'onnxruntime-gpu[cuda,cudnn]' numpy </dev/null >&2
  touch "$venv/.ready"
fi
command -v ffmpeg >/dev/null || { echo "ffmpeg was not found on $(hostname)" >&2; exit 1; }
work=$(mktemp -d)
trap 'rm -rf "$work"' EXIT
cat > "$work/worker.py" <<'TAKES_WORKER'
${worker}
TAKES_WORKER
cat > "$work/in.webm"
cd "$work"
"$venv/bin/python" worker.py "$1" </dev/null >&2
cat out.webm
`;
}

async function runRemote(job: Job, { src, out, W, H, ds, total }: Plan, host: string, setup: string): Promise<number> {
  const spec = { modelUrl: MODEL_URL, width: W, height: H, downsampleRatio: ds, decode: decoderArgs('in.webm', W, H), encode: encoderArgs('out.webm', W, H) };
  const command = `bash -lc ${shq(remoteScript(setup, await WORKER))} takes-matte ${shq(JSON.stringify(spec))}`;
  // BatchMode makes ssh fail instead of prompting for a password, since there is no terminal.
  const ssh = spawn('ssh', ['-T', '-o', 'BatchMode=yes', '--', host, command], { stdio: ['pipe', 'pipe', 'pipe'] });
  const closed = new Promise<number | null>((resolve, reject) => {
    ssh.on('error', reject);
    ssh.on('close', resolve);
  });
  const lines: string[] = []; // recent diagnostics, for the error message
  let n = 0;
  let partial = '';
  ssh.stderr.setEncoding('utf8');
  ssh.stderr.on('data', (d: string) => {
    const parts = (partial + d).split('\n');
    partial = parts.pop()!;
    for (const line of parts.map((l) => l.trimEnd())) {
      const m = /^takes:progress (\d+)$/.exec(line);
      if (m) {
        n = parseInt(m[1]);
        job.progress = Math.min(0.99, n / total);
      } else if (line.trim()) {
        console.log(`matte (${host}): ${line}`);
        lines.push(line);
        if (lines.length > 20) lines.shift();
      }
    }
  });
  try {
    const [code] = await Promise.all([
      closed,
      pipeline(createReadStream(src), ssh.stdin).catch(() => {}), // an early exit shows up in the exit code
      pipeline(ssh.stdout, createWriteStream(out)),
    ]);
    if (code !== 0) throw new Error(`Background removal on ${host} failed: ${lines.slice(-6).join('\n') || `ssh exited with code ${code}`}`);
    return n;
  } catch (e) {
    ssh.kill();
    throw e;
  }
}