import { InvocationType, InvokeCommand, LambdaClient, LogType } from '@aws-sdk/client-lambda';
import { spawn } from 'child_process';
import { writeFile } from 'fs';
import tmp from 'tmp';

export type Downbeat = { beat: number; time: number };

export const estimateBPM = (downbeats: Downbeat[], varianceThreshold: number = 0.05): number | null => {
  if (downbeats.length < 2) return null;

  let distances = downbeats.slice(1).map((downbeat, i) => downbeat.time - downbeats[i].time);
  const average = distances.reduce((a, b) => a + b, 0) / distances.length;
  const bpm = 60 / average;
  const variance = distances.reduce((a, b) => a + Math.abs(Math.log(b / average)), 0) / distances.length;
  return variance > varianceThreshold ? null : bpm;
};

const lambdaClient = new LambdaClient({ region: 'us-east-1' });

export const lambdaDetectDownbeats = async (sampleS3Key: string): Promise<Downbeat[]> => {
  const payload = { s3Key: sampleS3Key };
  const input = {
    FunctionName: 'DBNDownBeatTracker',
    InvocationType: InvocationType.RequestResponse,
    LogType: LogType.None,
    Payload: Buffer.from(JSON.stringify(payload), 'utf8'),
  };
  const res = await lambdaClient.send(new InvokeCommand(input));
  const json = JSON.parse(Buffer.from(res.Payload.buffer).toString('utf8'));
  if (!json.result) {
    throw new Error('Lambda invocation failed');
  }
  return json.result.map((x) => ({ time: x[0], beat: x[1] }));
};

export const detectDownbeats = async (buffer: Buffer): Promise<Downbeat[]> => {
  let cleanupCallback = () => {};
  try {
    const decodedFile = await new Promise<string>((resolve, reject) => {
      tmp.dir({ unsafeCleanup: true }, (err, path, innerCleanupCallback) => {
        if (err) {
          return reject(err);
        }
        cleanupCallback = innerCleanupCallback;

        const inputFilePath = `${path}/input`;
        const outputFilePath = `${path}/output.wav`;
        writeFile(inputFilePath, buffer, (err) => {
          if (err) {
            return reject(err);
          }
          const process = spawn('ffmpeg', [
            '-i',
            inputFilePath,
            '-ar',
            '44100',
            '-ac',
            '1',
            '-c:a',
            'pcm_s16le',
            '-f',
            'wav',
            outputFilePath,
          ]);
          process.on('error', (err) => reject(err));
          process.on('close', (code) => {
            if (code !== 0) {
              return reject(`ffmpeg process exited with code ${code}`);
            }
            resolve(outputFilePath);
          });
        });
      });
    });
    const madmomOutput = await new Promise<string>((resolve, reject) => {
      let output = '';
      const process = spawn('DBNDownBeatTracker', ['--norm', 'single', decodedFile]);
      process.on('error', (err) => reject(err));
      process.on('close', (code) => {
        if (code !== 0) {
          return reject(`madmom process exited with code ${code}`);
        }
        resolve(output);
      });
      process.stdout.on('data', (data) => {
        output += data.toString();
      });
    });
    return madmomOutput
      .split('\n')
      .filter((line) => line.trim().length > 0)
      .map((line) => {
        const [time, beat] = line.split(/\s+/);
        return { beat: parseInt(beat), time: parseFloat(time) };
      });
  } finally {
    cleanupCallback();
  }
};

export const detectDownbeatsAndBPM = async (
  {
    s3Key,
    lengthInMs,
  }: {
    s3Key: string;
    lengthInMs: number;
  },
  onProgress: (progress: number) => void
): Promise<{ bpm: number | null; downbeats: Downbeat[] }> => {
  const startTime = Date.now();

  // downbeat detection on lambda runs at 13.4x realtime
  let interval = setInterval(() => {
    const elapsed = (Date.now() - startTime) / 1000;
    if (elapsed > 0) {
      const estimatedProgress = (elapsed * 13.4) / (lengthInMs / 1000);
      onProgress(Math.tanh(1.5 * estimatedProgress));
    }
  }, 500);

  const downbeats = await (async () => {
    try {
      return await lambdaDetectDownbeats(s3Key);
    } finally {
      clearInterval(interval);
    }
  })();

  onProgress(1);

  return {
    bpm: estimateBPM(downbeats),
    downbeats,
  };
};