import { Socket } from 'socket.io';
import { v4 as uuidv4 } from 'uuid';
import { getRedisClient, withIsolatedRedisClient } from '../external-services/redis.js';
import { SocketEvent } from '../server-utils/socketMessageHandler.js';
import { getSample, getSignedUrlForSampleS3Key, insertSample } from './samples.js';
import { getUserById } from './users.js';

// TODO: share type with frontend
enum AudiocraftAction {
  GenerateAudio = 'generateAudio',
}

export const audiocraftMessageHandler = async (socket: Socket, message: any) => {
  const { textPrompt, audioPromptSampleUuid, duration, codec, bpm } = message;
  const action = message.action as AudiocraftAction;
  if (action === AudiocraftAction.GenerateAudio) {
    socket.emit(SocketEvent.Final, {
      url: await audiocraftGenerateAudio(
        textPrompt,
        textPrompt,
        audioPromptSampleUuid,
        duration || 10,
        codec || 'mp3',
        (socket.request as any).user.id,
        bpm,
        (p) => socket.emit(SocketEvent.Progress, p)
      ),
    });
    return;
  }
};

export const audiocraftGenerateAudio = async (
  textPrompt: string,
  sampleName: string,
  audioPromptSampleUuid: string,
  duration: number,
  codec: string,
  userId: number,
  bpm: number | undefined,
  progressCallback: (p: number) => void,
  multiple?: boolean
) => {
  const requestId = uuidv4();
  const queueName = 'musicgen_requests_beta_2';
  const timeout = 120;

  let audioPromptS3Key = null;
  if (audioPromptSampleUuid) {
    audioPromptS3Key = (
      await getSample({ uuid: audioPromptSampleUuid, asUser: await getUserById(userId) }, { returnCompressed: true })
    ).s3Key;
  }

  const redisClient = getRedisClient();
  const now = Date.now() / 1000.0;
  await redisClient.lpush(
    queueName,
    JSON.stringify({
      id: requestId,
      data: {
        textPrompt,
        audioPromptS3Key,
        duration,
        codec,
        bpm,
      },
      timeSubmitted: now,
    })
  );
  const res = await withIsolatedRedisClient(timeout * 1000 + 2000, async (isolatedClientPromise) => {
    const isolatedClient = await isolatedClientPromise;
    while (true) {
      const res = await isolatedClient.blpop(requestId, timeout);
      if (!res) {
        throw new Error('Timeout waiting for response');
      }
      const message = JSON.parse(res[1]);
      if (message.progress) {
        progressCallback(message.progress);
      } else {
        return message;
      }
    }
  });
  if (res.error) {
    throw new Error(res.error);
  }

  const { s3Key, results } = res;

  if (multiple) {
    const output: { uuid: string; url: string }[] = [];
    for (let i = 0; i < results.length; i++) {
      const insertedRow = await insertSample({
        name: sampleName,
        uuid: results[i].s3Key.split('.')[0],
        s3Key: results[i].s3Key,
        authorId: userId,
      });

      output.push({
        uuid: insertedRow.uuid,
        url: await getSignedUrlForSampleS3Key(results[i].s3Key),
      });
    }

    return output;
  } else {
    const insertedRow = await insertSample({
      name: sampleName,
      uuid: s3Key.split('.')[0],
      s3Key,
      authorId: userId,
    });

    return {
      uuid: insertedRow.uuid,
      url: await getSignedUrlForSampleS3Key(s3Key),
    };
  }
};