#!/usr/bin/env node
import { resolve } from "path";
import { program } from "commander";
import { createReadStream, createWriteStream, ReadStream } from "fs";
import { readFile } from "fs/promises";
import {
  initDspModule,
  streamPCMChunks,
  getStateFileRenderLengthSeconds,
} from "./renderer.js";
import { BufferF32 } from "../../artifacts/dsp-engine.js";

const SAMPLE_RATE = 48_000;

interface CliOptions {
  input?: string;
  output?: string;
  start?: number;
  end?: number;
  getRenderLength: boolean;
  bufferSize: number;
  mediaMap?: string;
  disableLimiter: undefined | true;
  streamPcmPipe?: string;
  streamPcmCopyPipe?: string;
  streamSum: string[];
  frames?: number;
  onlyTrackId?: string;
}

const parseIntSafe = (value: string) => {
  const int = Number.parseInt(value, 10);
  if (Number.isNaN(int)) {
    throw new Error(`Invalid integer: ${value}`);
  }
  return int;
};

const parseFiniteFloat = (value: string) => {
  const float = Number.parseFloat(value);
  if (Number.isNaN(float) || !Number.isFinite(float)) {
    throw new Error(`Invalid float: ${value}`);
  }
  return float;
};

const copyTypedArrayToBuffer = (typedArray: Float32Array | Float64Array) => {
  return Buffer.from(
    new Uint8Array(
      typedArray.buffer,
      typedArray.byteOffset,
      typedArray.byteLength
    )
  );
};

const loadMediaMap = async (mediaMapPath: string) => {
  const mediaMap = JSON.parse(
    await readFile(resolve(mediaMapPath), "utf-8")
  ) as Record<string, string>;
  return (clipId: string) => mediaMap[clipId];
};

const handleStreamSum = async (options: CliOptions) => {
  const { dspModule, dspContext } = await initDspModule(SAMPLE_RATE, true);
  const bouncer = dspContext.createSumOnlyBouncer(options.bufferSize);
  const planarWorkBuffer = new dspModule.BufferF32(2, options.bufferSize);
  const interleavedWorkBuffer = new dspModule.BufferF32(
    1,
    2 * options.bufferSize
  );
  const needBytes = options.bufferSize * 4 * 2;
  const streams = Array.from({ length: options.streamSum.length }, () => ({
    buffers: [] as Buffer[],
    isFinished: false,
    readStream: null as ReadStream | null,
  }));
  const writer = createWriteStream(options.streamPcmPipe!, {
    highWaterMark: 32768,
  });
  let waitingForDrain = false;
  let framesSubmitted = 0;
  let framesWritten = 0;
  const takeResult = (result: BufferF32 | null) => {
    try {
      if (result && result.frameCount > 0) {
        const framesLeftToWrite =
          options.frames !== undefined
            ? options.frames - framesWritten
            : undefined;
        const tempInterleaved = new dspModule.BufferF32(
          1,
          2 * result.frameCount
        );
        try {
          result.interleaveTo(tempInterleaved);
          const framesToWrite = Math.max(
            0,
            Math.min(
              result.frameCount,
              framesLeftToWrite === undefined
                ? result.frameCount
                : framesLeftToWrite
            )
          );
          framesWritten += framesToWrite;
          return writer.write(
            copyTypedArrayToBuffer(
              tempInterleaved.view()[0].subarray(0, 2 * framesToWrite)
            )
          );
        } finally {
          tempInterleaved.delete();
        }
      }
      return true;
    } finally {
      result?.delete();
    }
  };
  const progress = () => {
    if (waitingForDrain) {
      return;
    }
    const isFinished = streams.every((stream) => stream.isFinished);
    const canSubmitFrame = streams.every(
      (stream) =>
        stream.buffers.reduce((acc, buf) => acc + buf.length, 0) >= needBytes
    );
    const framesRemaining =
      options.frames !== undefined
        ? options.frames - framesSubmitted
        : undefined;
    let needsWriterDrain = false;

    if (canSubmitFrame) {
      for (const stream of streams) {
        const combinedBuffer = Buffer.concat(stream.buffers);
        const prefix = combinedBuffer.subarray(0, needBytes);
        const suffix = combinedBuffer.subarray(needBytes);
        stream.buffers = [suffix];
        const pfxa = new Float32Array(
          prefix.buffer,
          prefix.byteOffset,
          prefix.byteLength / Float32Array.BYTES_PER_ELEMENT
        );
        interleavedWorkBuffer.set(0, 0, pfxa);
        planarWorkBuffer.deinterleaveFrom(interleavedWorkBuffer);
        bouncer.sumIntoMixdownBuffer(planarWorkBuffer);
      }
      needsWriterDrain ||= !takeResult(bouncer.bounceNext(false));
      framesSubmitted += options.bufferSize;
    } else if (
      isFinished &&
      framesRemaining !== undefined &&
      framesRemaining > 0
    ) {
      needsWriterDrain ||= !takeResult(bouncer.bounceNext(false));
      framesSubmitted += options.bufferSize;
    }

    // allow write stream to catch up
    if (needsWriterDrain) {
      for (const stream of streams) {
        stream.readStream?.pause();
      }
      waitingForDrain = true;
      writer.once("drain", () => {
        waitingForDrain = false;
        progress();
      });
      return;
    }

    if (!isFinished) {
      // update backpressure on readers
      let isSomeReaderRunning = false;
      for (const stream of streams) {
        if (stream.isFinished) {
          continue;
        }
        if (
          stream.buffers.reduce((acc, buf) => acc + buf.length, 0) >= needBytes
        ) {
          stream.readStream?.pause();
        } else {
          isSomeReaderRunning = true;
          stream.readStream?.resume();
        }
      }
      if (!isSomeReaderRunning) {
        // event loop will be empty - we could quit too early
        if (
          !streams.every(
            (stream) =>
              stream.buffers.reduce((acc, buf) => acc + buf.length, 0) >=
              needBytes
          )
        ) {
          console.error(
            "some pipes contained different amounts of data! aborting"
          );
          process.exit(1);
        }

        return progress();
      }
    } else if (canSubmitFrame) {
      // all streams are finished but there are more frames to submit, go again
      return progress();
    } else {
      // all streams are finished and all frames are submitted, drain then we're done
      while (true) {
        const result = bouncer.bounceNext(true);
        if (!result || result.frameCount === 0) {
          break;
        }
        takeResult(result);
      }
      writer.end();
    }
  };
  for (const [i, path] of options.streamSum.entries()) {
    const readStream = createReadStream(path, {
      highWaterMark: 32768,
    });
    streams[i].readStream = readStream;
    readStream.on("error", (error) => {
      console.error(
        `Error reading from stream index ${i} path ${path}:`,
        error
      );
      process.exit(1);
    });
    readStream.on("data", (data) => {
      if (!(data instanceof Buffer)) {
        throw new Error(`Expected Buffer, got ${typeof data}`);
      }
      streams[i].buffers.push(data);
      progress();
    });
    readStream.on("end", () => {
      streams[i].isFinished = true;
      progress();
    });
  }
};

const handleStreamPcmPipe = async (options: CliOptions) => {
  const writer = createWriteStream(options.streamPcmPipe!, {
    highWaterMark: 32768,
  });
  const copyWriter = options.streamPcmCopyPipe
    ? createWriteStream(options.streamPcmCopyPipe!, {
        highWaterMark: 32768,
      })
    : null;
  try {
    let totalFrames = 0;
    for await (const { interleaved, frameCount } of streamPCMChunks(
      resolve(options.input!),
      options.start!,
      options.end ?? +Infinity,
      SAMPLE_RATE,
      options.bufferSize,
      !options.disableLimiter,
      await loadMediaMap(options.mediaMap!),
      options.onlyTrackId
    )) {
      const framesToWrite = Math.max(
        0,
        Math.min(
          frameCount,
          options.frames !== undefined
            ? options.frames - totalFrames
            : frameCount
        )
      );
      if (framesToWrite === 0) {
        break;
      }
      totalFrames += framesToWrite;
      // Stream raw f32le (interleaved stereo)
      const interleavedPrefix = interleaved.subarray(0, 2 * framesToWrite);
      const interleavedPrefixCopy = copyTypedArrayToBuffer(interleavedPrefix);
      const written = writer.write(interleavedPrefixCopy);
      if (!written) {
        await new Promise<void>((resolve) => writer.once("drain", resolve));
      }
      if (copyWriter) {
        const copyWritten = copyWriter.write(interleavedPrefixCopy);
        if (!copyWritten) {
          await new Promise<void>((resolve) =>
            copyWriter.once("drain", resolve)
          );
        }
      }
    }
    const durationSec = totalFrames / SAMPLE_RATE;
    console.error(
      `streamed ${totalFrames} frames (${durationSec.toFixed(2)}s)`
    );
  } finally {
    writer.end();
    if (copyWriter) {
      copyWriter.end();
    }
  }
};

const handleNormalRender = async (options: CliOptions) => {
  const { tmpdir } = await import("os");
  const { mkdtemp, rm } = await import("fs/promises");
  const { join } = await import("path");
  const { execFile } = await import("child_process");
  const tmpDir = await mkdtemp(join(tmpdir(), "bounce-tmp-"));
  const outputPcmFile = join(tmpDir, "output.pcm"); // just a normal file to accumulate the PCM data
  try {
    await handleStreamPcmPipe({
      ...options,
      streamPcmPipe: outputPcmFile,
    });
    await new Promise<void>((resolve, reject) => {
      execFile(
        "ffmpeg",
        [
          "-hide_banner",
          "-loglevel",
          "error",
          "-y",
          "-f",
          "f32le",
          "-ar",
          `${SAMPLE_RATE}`,
          "-ac",
          "2",
          "-i",
          outputPcmFile,
          options.output!,
        ],
        (error, stdout, stderr) => {
          if (error) {
            console.error(stderr);
            reject(error);
          } else {
            resolve();
          }
        }
      );
    });
  } finally {
    await rm(tmpDir, { recursive: true, force: true });
  }
};

program
  .name("bounce")
  .description(
    "Command line utility to render studio session file to WAV audio"
  )
  .version("1.0.0");

const collectSumPath = (value: string, previous: string[]) => {
  return previous.concat([value]);
};

program
  .option("-i, --input <statefile>", "Path to the studio state JSON file")
  .option("-o, --output <path>", "Output file path (default: output.wav)")
  .option("-s, --start <beats>", "Start position in beats", parseFiniteFloat)
  .option("-e, --end <beats>", "End position in beats", parseFiniteFloat)
  .option(
    "-f, --frames <frames>",
    "Render exactly this many frames",
    parseIntSafe
  )
  .option("-m, --media-map <path>", "Map of clip IDs to media paths")
  .option(
    "-b, --buffer-size <size>",
    "Buffer size for bounce (default: 4096)",
    parseIntSafe,
    4096
  )
  .option(
    "--stream-pcm-pipe <path>",
    "Stream interleaved PCM to a named pipe (chunked)"
  )
  .option(
    "--stream-pcm-copy-pipe <path>",
    "Copy interleaved PCM to a named pipe (chunked)"
  )
  .option(
    "-u, --stream-sum <path>",
    "Sum over a number of input PCM pipes into a single output pipe with limiting",
    collectSumPath,
    []
  )
  .option("--disable-limiter", "Disable limiter (default: enabled)")
  .option("--get-render-length", "Get render length in frames and exit")
  .option("--only-track-id <id>", "Only render the track with the given ID")
  .action(async (options: CliOptions) => {
    if (options.start !== undefined && options.end !== undefined) {
      if (options.start >= options.end) {
        console.error("--start must be less than --end");
        process.exit(2);
      }
    }
    if (options.bufferSize < 1 || options.bufferSize > 16384) {
      console.error("--buffer-size must be between 1 and 16384");
      process.exit(2);
    }
    if (options.frames !== undefined) {
      if (options.end !== undefined) {
        console.error("--frames and --end are not supported together");
        process.exit(2);
      }
      if (options.frames <= 0) {
        console.error("--frames must be greater than 0");
        process.exit(2);
      }
    }
    try {
      if (options.getRenderLength) {
        for (const requiredOption of ["input", "start", "end"]) {
          if (options[requiredOption as keyof CliOptions] === undefined) {
            console.error(
              `${requiredOption} is required for --get-render-length`
            );
            process.exit(2);
          }
        }
        for (const incompatibleOption of [
          "output",
          "mediaMap",
          "disableLimiter",
          "streamPcmPipe",
          "streamPcmCopyPipe",
          "frames",
          "onlyTrackId",
        ]) {
          if (options[incompatibleOption as keyof CliOptions] !== undefined) {
            console.error(
              `${incompatibleOption} is not supported with --get-render-length`
            );
            process.exit(2);
          }
        }
        if (options.streamSum.length > 0) {
          console.error(
            "--stream-sum is not supported with --get-render-length"
          );
          process.exit(2);
        }

        const renderLength = await getStateFileRenderLengthSeconds(
          options.input!,
          options.start!,
          options.end!
        );
        console.log(
          `Render length: ${Math.ceil(renderLength * SAMPLE_RATE)} frames`
        );
        process.exit(0);
      } else if (options.streamSum.length > 0) {
        // stream sum mode
        for (const incompatibleOption of [
          "input",
          "output",
          "mediaMap",
          "start",
          "end",
          "disableLimiter",
          "streamPcmCopyPipe",
          "onlyTrackId",
        ]) {
          if (options[incompatibleOption as keyof CliOptions] !== undefined) {
            console.error(
              `${incompatibleOption} is not supported with --stream-sum`
            );
            process.exit(2);
          }
        }

        if (!options.streamPcmPipe) {
          console.error("--stream-sum requires --stream-pcm-pipe <path>");
          process.exit(2);
        }

        await handleStreamSum(options);
      } else if (options.streamPcmPipe) {
        // normal input, stream to a named pipe
        if (options.output) {
          console.error(
            "--stream-pcm-pipe and --output are not supported together"
          );
          process.exit(2);
        }
        for (const requiredOption of ["input", "mediaMap", "start"]) {
          if (options[requiredOption as keyof CliOptions] === undefined) {
            console.error(
              `${requiredOption} is required for --stream-pcm-pipe`
            );
            process.exit(2);
          }
        }
        if (options.end === undefined && options.frames === undefined) {
          console.error("--end or --frames is required for --stream-pcm-pipe");
          process.exit(2);
        }

        await handleStreamPcmPipe(options);
      } else {
        // normal input, write to a file
        for (const requiredOption of ["input", "output", "mediaMap", "start"]) {
          if (options[requiredOption as keyof CliOptions] === undefined) {
            console.error(`${requiredOption} is required for normal render`);
            process.exit(2);
          }
        }
        if (options.streamPcmCopyPipe) {
          console.error(
            "--stream-pcm-copy-pipe is not supported with normal render"
          );
          process.exit(2);
        }
        if (options.end === undefined && options.frames === undefined) {
          console.error("--end or --frames is required for normal render");
          process.exit(2);
        }
        await handleNormalRender(options);
      }
    } catch (error) {
      console.error("Error:", error);
      process.exit(1);
    }
  });

program.parse();