import { readFile, open as openFile } from "fs/promises";
import { readSync, fstatSync } from "fs";
import type {
  MainModule,
  DSPContext,
  FfmpegAudioBuffer,
} from "../../../artifacts/dsp-engine.js";
import createDspEngine from "../../../artifacts/dsp-engine.js";
import { promisify } from "util";
import { execFile, spawn } from "child_process";

const BUFFER_SIZE = 4096;

const execFileAsync = promisify(execFile);

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

export async function initDspModule(
  channelCount: number,
  sampleRate: number
): Promise<{
  dspModule: MainModule;
  dspContext: DSPContext;
}> {
  const dspModule = await (
    createDspEngine as unknown as () => Promise<MainModule>
  )();
  const dspContext = await dspModule.DSPContext.createOffline(
    "bounce",
    0.5,
    channelCount,
    sampleRate,
    false
  );

  return { dspModule, dspContext };
}

export async function getAudioFileFormatInfo(
  inputFilePath: string
): Promise<{ channelCount: number; sampleRate: number }> {
  try {
    const result = await execFileAsync("ffprobe", [
      "-v",
      "error",
      "-select_streams",
      "a:0",
      "-show_entries",
      "stream=sample_rate,channels",
      "-of",
      "json",
      inputFilePath,
    ]);
    const json = JSON.parse(result.stdout.toString());
    const stream = json.streams?.[0];
    if (!stream || !stream.sample_rate || !stream.channels) {
      throw new Error("Failed to get audio format info from ffprobe output");
    }
    return {
      channelCount: parseInt(stream.channels),
      sampleRate: parseInt(stream.sample_rate),
    };
  } catch (e) {
    throw new Error(`Failed to get audio format info: ${e}`);
  }
}

function execFromBuffers(
  command: string,
  args: string[]
): { send: (buffer: Buffer) => Promise<void>; close: () => Promise<void> } {
  const childProcess = spawn(command, args, {
    stdio: ["pipe", "ignore", "inherit"],
  });

  const send = async (buffer: Buffer) => {
    return new Promise<void>((resolve, reject) => {
      childProcess.stdin?.write(buffer, (error) => {
        if (error) {
          reject(error);
        } else {
          resolve();
        }
      });
    });
  };

  const close = async () => {
    childProcess.stdin?.end();

    return await new Promise<void>((resolve, reject) => {
      childProcess.on("close", (code) => {
        if (code === 0) {
          resolve();
        } else {
          reject(new Error(`child process exited with code ${code}`));
        }
      });
      childProcess.on("error", (error) => {
        reject(error);
      });
    });
  };

  return { send, close };
}

export async function openAudioFile(
  inputFilePath: string,
  dspModule: MainModule
): Promise<FfmpegAudioBuffer> {
  const handle = await openFile(inputFilePath, "r");
  try {
    const size = fstatSync(handle.fd).size;
    return dspModule.FfmpegAudioBuffer.createFromReadCallback(
      size,
      (buf: Uint8Array, offset: number) => {
        const bytesRed = readSync(handle.fd, buf, 0, buf.length, offset);
        if (bytesRed !== buf.length) {
          // this should never happen
          console.error(
            `Short read from ${inputFilePath} at offset ${offset}: got ${bytesRed} bytes, expected ${buf.length}`
          );
        }
      },
      () => {
        (async () => {
          // be sure to retain handle until it's closed
          await handle.close();
        })();
      }
    );
  } catch (error) {
    await handle.close();
    throw error;
  }
}

export function writeAudioBuffersToFile(
  outputFilePath: string,
  channelCount: number,
  sampleRate: number
) {
  return execFromBuffers("ffmpeg", [
    "-v",
    "error",
    "-f",
    "f32le",
    "-ar",
    sampleRate.toString(),
    "-ac",
    channelCount.toString(),
    "-y",
    "-i",
    "-",
    outputFilePath,
  ]);
}

interface WarpMarker {
  timeInUnderlyingBuffer: number;
  timeInOutput: number;
}

async function loadWarpMarkers(
  warpMarkersFilePath: string
): Promise<WarpMarker[]> {
  const warpMarkers = await readFile(warpMarkersFilePath, "utf8");
  return warpMarkers
    .split("\n")
    .filter((line) => line.trim())
    .map((line) => {
      const [timeInUnderlyingBuffer, timeInOutput] = line
        .trim()
        .split(" ")
        .map(Number);
      return { timeInUnderlyingBuffer, timeInOutput };
    })
    .sort((a, b) => a.timeInOutput - b.timeInOutput)
    .filter(
      (marker, index, array) =>
        index === 0 || marker.timeInOutput !== array[index - 1].timeInOutput
    );
}

export async function timestretch(
  inputFilePath: string,
  outputFilePath: string,
  transpose: number,
  warpMarkersFilePath: string | null
): Promise<void> {
  const { channelCount, sampleRate } = await getAudioFileFormatInfo(
    inputFilePath
  );
  if (!(channelCount == 1 || channelCount == 2)) {
    throw new Error(`Unsupported channel count: ${channelCount}`);
  }
  const { dspModule, dspContext } = await initDspModule(
    channelCount,
    sampleRate
  );
  const ffmpegAudioBuffer = await openAudioFile(inputFilePath, dspModule);
  try {
    const warpMarkers: WarpMarker[] = warpMarkersFilePath
      ? await loadWarpMarkers(warpMarkersFilePath)
      : [];
    const inputDuration = ffmpegAudioBuffer.frameCount / sampleRate;
    console.log(`Input duration: ${inputDuration} seconds`);
    if (transpose < -24 || transpose > 24) {
      throw new Error(
        `Transpose must be between -24 and 24 semitones: ${transpose}`
      );
    }
    for (const warpMarker of warpMarkers) {
      if (
        warpMarker.timeInUnderlyingBuffer < 0 ||
        warpMarker.timeInUnderlyingBuffer > inputDuration
      ) {
        throw new Error(
          `Warp marker out of bounds: ${warpMarker.timeInUnderlyingBuffer} (must be between 0 and ${inputDuration} seconds)`
        );
      }
    }
    let outputDuration: number;
    if (warpMarkers.length == 0) {
      outputDuration = inputDuration;
    } else if (warpMarkers.length == 1) {
      outputDuration =
        inputDuration +
        warpMarkers[0].timeInOutput -
        warpMarkers[0].timeInUnderlyingBuffer;
    } else {
      const lastSegmentRate =
        (warpMarkers[warpMarkers.length - 1].timeInUnderlyingBuffer -
          warpMarkers[warpMarkers.length - 2].timeInUnderlyingBuffer) /
        (warpMarkers[warpMarkers.length - 1].timeInOutput -
          warpMarkers[warpMarkers.length - 2].timeInOutput);
      if (lastSegmentRate == 0.0) {
        // not well defined, just use input duration
        outputDuration = inputDuration;
      } else {
        const inputTimeRemaining =
          lastSegmentRate > 0.0
            ? // going forwards to end
              inputDuration -
              warpMarkers[warpMarkers.length - 1].timeInUnderlyingBuffer
            : // going backwards to start
              warpMarkers[warpMarkers.length - 1].timeInUnderlyingBuffer;
        outputDuration =
          warpMarkers[warpMarkers.length - 1].timeInOutput +
          inputTimeRemaining / Math.abs(lastSegmentRate);
      }
    }
    const timing = dspModule.TimelineTempoMap.fromArray(1, []);
    const warpMap = dspModule.WarpMap.fromArray(warpMarkers);
    const audioClip = dspContext.createAudioClip(
      ffmpegAudioBuffer,
      warpMap,
      1.0, // gain
      0.0, // timelineStartBeats
      outputDuration, // timelineEndBeats
      0.0, // loopStartBeats
      outputDuration, // loopEndBeats
      0.0, // readStartBeats
      0.0, // fadeInBeats
      1.0, // fadeInExponent
      0.0, // fadeOutBeats
      1.0, // fadeOutExponent
      transpose, // transposition
      1.0, // warpedContentBps
      true, // warpEnabled
      "" // clipId
    );
    const track = dspContext.createTrack(1.0, 0.0, false, null, null, [
      audioClip,
    ]);
    const timeline = dspContext.createTimeline(timing, [track]);

    timeline.fadeInStartBeats = -Infinity;
    timeline.fadeInLengthBeats = 0;
    timeline.fadeInExponent = 0;
    timeline.fadeOutEndBeats = Infinity;
    timeline.fadeOutLengthBeats = 0;
    timeline.fadeOutExponent = 0;
    timeline.masterGain = 1;

    dspContext.swapLiveTimeline(timeline);
    // release local handles
    timing.delete();
    warpMap.delete();
    audioClip.delete();
    track.delete();
    timeline.delete();
    // bounce

    const { send: writeAudioBuffer, close: closeAudioFile } =
      writeAudioBuffersToFile(outputFilePath, channelCount, sampleRate);

    const outputLength = outputDuration * sampleRate;
    let outputOffset = 0;

    const bouncer = dspContext.createBouncer(0.0, outputDuration, BUFFER_SIZE);
    const interleavedOutputBuffer = new dspModule.BufferF32(
      1,
      channelCount * BUFFER_SIZE
    );

    try {
      while (true) {
        const outputBuffer = bouncer.bounceNext();
        if (outputBuffer == null) {
          break;
        }
        try {
          outputBuffer.interleaveTo(interleavedOutputBuffer);
          const frameCount = outputBuffer.frameCount;
          const endOfBuffer =
            Math.min(outputOffset + frameCount, outputLength) - outputOffset;
          if (endOfBuffer <= 0) {
            return;
          }
          outputOffset += endOfBuffer;
          await writeAudioBuffer(
            copyTypedArrayToBuffer(
              interleavedOutputBuffer
                .view()[0]
                .subarray(0, endOfBuffer * channelCount)
            )
          );
        } finally {
          outputBuffer.delete();
        }
      }
    } finally {
      await closeAudioFile();

      bouncer.delete();
      dspContext.delete();
      interleavedOutputBuffer.delete();
    }
  } finally {
    ffmpegAudioBuffer.delete();
  }
}