import { spawn, type ChildProcessByStdio } from "node:child_process";
import type { Readable, Writable } from "node:stream";
import { mkdir, rm } from "node:fs/promises";
import { join } from "node:path";
import { CoreError, fail } from "./errors.ts";

export type FrameWriter = (frame: number, png: Uint8Array) => Promise<void>;

export interface FrameStreamOptions {
  directory: string;
  frameCount: number;
  parallelism: number;
  fpsSource: string;
  output: string;
  ffmpegPath?: string;
  signal?: AbortSignal;
  acquireEncoder?: (signal: AbortSignal) => Promise<() => void>;
  onEncoder?: (details: { pid: number | undefined; startFrame: number; endFrame: number }) => void;
  onComplete?: (details: { frames: number; encoders: number }) => void;
}

async function waitEncoder(child: ChildProcessByStdio<Writable | null, null, Readable>): Promise<void> {
  let launchError: Error | undefined;
  let readError: unknown;
  let message = "";
  child.once("error", (error) => { launchError = error; });
  const exited = new Promise<number | null>((resolve) => child.once("close", resolve));
  child.stderr.on("data", (chunk) => { message = (message + String(chunk)).slice(-12_000); });
  child.stderr.once("error", (error) => { readError = error; });
  const code = await exited;
  if (launchError !== undefined) fail("FFMPEG_START_FAILED", "无法启动视觉编码器", { reason: launchError.message });
  if (readError !== undefined) throw readError;
  if (code !== 0) fail("VISUAL_CACHE_ENCODE_FAILED", "视觉流式编码失败", { exitCode: code, stderr: message });
}

/** PNGs go straight from the sampler into independent lossless encoders. */
export async function withFrameStream<T>(
  options: FrameStreamOptions,
  render: (write: FrameWriter) => Promise<T>,
): Promise<{ value: T; mediaPath?: string }> {
  if (!Number.isSafeInteger(options.frameCount) || options.frameCount < 1 ||
      !Number.isSafeInteger(options.parallelism) || options.parallelism < 1) {
    fail("INVALID_RENDER_OPTION", "流式编码的帧数和并发数必须是正整数");
  }
  const directory = join(options.directory, ".stream");
  await mkdir(directory, { recursive: true });
  const count = Math.min(options.parallelism, options.frameCount);
  const size = Math.floor(options.frameCount / count);
  const extra = options.frameCount % count;
  const ffmpeg = options.ffmpegPath ?? "ffmpeg";
  type Child = ChildProcessByStdio<Writable, null, Readable>;
  type Pending = { png: Uint8Array; resolve: () => void; reject: (error: unknown) => void };
  const parts = Array.from({ length: count }, (_, index) => {
    const start = index * size + Math.min(index, extra);
    return { start, end: start + size + (index < extra ? 1 : 0), next: start,
      path: join(directory, `${index}.mov`), pending: new Map<number, Pending>(),
      pumping: undefined as Promise<void> | undefined, child: undefined as Child | undefined, finished: undefined as Promise<void> | undefined };
  });
  let failure: unknown;
  let failed = false;
  let written = 0;
  const controller = new AbortController();
  let concat: ChildProcessByStdio<null, null, Readable> | undefined;
  let concatenated: Promise<void> | undefined;
  const stop = (error: unknown) => {
    if (!failed) { failed = true; failure = error; }
    controller.abort();
    for (const part of parts) {
      for (const item of part.pending.values()) item.reject(failure);
      part.pending.clear();
      if (part.child !== undefined) {
        // FFmpeg can remain blocked reading stdin after SIGTERM. Close the
        // producer and terminate this disposable staging encoder immediately.
        part.child.stdin.destroy();
        if (part.child.exitCode === null) { try { part.child.kill("SIGKILL"); } catch { /* exited */ } }
      }
    }
    if (concat?.exitCode === null) { try { concat.kill("SIGKILL"); } catch { /* exited */ } }
  };
  const abort = () => stop(new CoreError("RENDER_CANCELLED", "视觉流式编码已取消"));
  options.signal?.addEventListener("abort", abort, { once: true });
  if (options.signal?.aborted) abort();
  const pump = (part: typeof parts[number]) => {
    if (part.pumping !== undefined) return;
    part.pumping = (async () => {
      try {
        while (!failed) {
          const item = part.pending.get(part.next);
          if (item === undefined) return;
          if (part.child === undefined) {
            const release = await options.acquireEncoder?.(controller.signal) ?? (() => {});
            if (failed) { release(); throw failure; }
            try {
              part.child = spawn(ffmpeg, [ "-hide_banner", "-loglevel", "error", "-y",
                "-threads", "1", "-f", "image2pipe", "-framerate", options.fpsSource, "-c:v", "png", "-i", "pipe:0",
                "-filter_threads", "1", "-frames:v", String(part.end - part.start), "-an", "-c:v", "qtrle", "-pix_fmt", "argb", part.path],
              { stdio: ["pipe", "ignore", "pipe"] });
            } catch (error) { release(); throw error; }
            const child = part.child;
            part.finished = waitEncoder(child).finally(release);
            child.stdin.on("error", stop);
            void part.finished.catch(stop);
            options.onEncoder?.({ pid: child.pid, startFrame: part.start, endFrame: part.end });
          }
          const child = part.child;
          await new Promise<void>((resolve, reject) => child.stdin.write(item.png, (error) => error ? reject(error) : resolve()));
          if (failed) throw failure;
          part.pending.delete(part.next++);
          written++;
          item.resolve();
          if (part.next === part.end) part.child.stdin.end();
        }
      } catch (error) { stop(error); }
    })().finally(() => {
      part.pumping = undefined;
      if (!failed && part.pending.has(part.next)) pump(part);
    });
  };
  const write: FrameWriter = (frame, png) => {
    if (failed) return Promise.reject(failure);
    const part = parts.find((part) => frame >= part.start && frame < part.end);
    if (!Number.isSafeInteger(frame) || part === undefined || frame < part.next || part.pending.has(frame)) {
      return Promise.reject(new CoreError("INTERNAL_ERROR", "视觉帧编号重复或越界", { frame }));
    }
    const completed = new Promise<void>((resolve, reject) => part.pending.set(frame, { png, resolve, reject }));
    pump(part);
    return completed;
  };
  try {
    if (failed) throw failure;
    const value = await render(write);
    if (failed) throw failure;
    // A static artifact writes its one PNG through the existing static path.
    if (written === 0) return { value };
    if (written !== options.frameCount) fail("INTERNAL_ERROR", "视觉流缺少帧", { written, expected: options.frameCount });
    await Promise.all(parts.map((part) => part.finished));
    if (failed) throw failure;
    const list = join(directory, "concat.txt");
    await Bun.write(list, "ffconcat version 1.0\n" + parts.map((part) => `file '${part.path.replace(/'/g, "'\\''")}'`).join("\n") + "\n");
    const mediaPath = options.output;
    concat = spawn(ffmpeg, [ "-hide_banner", "-loglevel", "error", "-y", "-f", "concat", "-safe", "0", "-i", list,
      "-c", "copy", mediaPath], { stdio: ["ignore", "ignore", "pipe"] });
    concatenated = waitEncoder(concat);
    await concatenated;
    if (failed) throw failure;
    options.onComplete?.({ frames: written, encoders: count });
    return { value, mediaPath };
  } catch (error) {
    stop(error);
    throw error;
  } finally {
    options.signal?.removeEventListener("abort", abort);
    // Producers and processes must finish before staging can be removed.
    await Promise.allSettled(parts.flatMap((part) => part.pumping === undefined ? [] : [part.pumping]));
    await Promise.allSettled(parts.flatMap((part) => part.finished === undefined ? [] : [part.finished]));
    if (concatenated !== undefined) await concatenated.catch(() => undefined);
    await rm(directory, { recursive: true, force: true });
  }
}