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 { encode } from '../signal-processing/decode.js';
import decode from '../signal-processing/decodeAudio.js';
import { User } from '../types/serverTypes.js';
import runServiceWithQuota from '../utils/runServiceWithQuota.js';
import { putObject } from '../utils/s3Buffers.js';
import { SAMPLES_BUCKET, getSample, getSignedUrlForSampleS3Key, insertSample } from './samples.js';

// a minus b
export const subtract = async (uuidA: string, uuidB: string, user: User) => {
  const cacheKey = `${uuidA} minus ${uuidB}`;
  const cachedSample = await getSample({ cacheKey, asUser: user }, { returnCompressed: false });

  if (cachedSample) {
    return cachedSample.uuid;
  }

  const sampleA = await getSample({ uuid: uuidA, asUser: user }, { returnCompressed: false, returnBuffer: true });
  const sampleB = await getSample({ uuid: uuidB, asUser: user }, { returnCompressed: false, returnBuffer: true });

  const bufferA = await decode(sampleA.buffer, { sampleRate: 44100 });
  const bufferB = await decode(sampleB.buffer, { sampleRate: 44100 });

  const result = {
    sampleRate: bufferA.sampleRate,
    channelData: bufferA.channelData.map((channel, i) =>
      channel.map((sample, j) => sample - bufferB.channelData[i][j])
    ),
  };

  const sampleUUID = uuidv4();
  const s3Key = sampleUUID + '.wav';
  await putObject(
    SAMPLES_BUCKET,
    s3Key,
    await encode(result.channelData, { sampleRate: result.sampleRate, bitDepth: 16 })
  );

  const insertedSample = await insertSample({
    name: `Subtract ${uuidA} - ${uuidB}`,
    uuid: sampleUUID,
    s3Key: s3Key,
    authorId: user.id,
    warp: sampleA.warp,
    cacheKey,
    derivedFromSampleId: sampleA.id,
    derivationType: 'subtract',
    tags: [],
  });

  return insertedSample.uuid;
};

export const runDemucs = async (
  user: User,
  uuid: string,
  samplePackId?: number,
  onProgress?: (progressMessage) => void
): Promise<{ [key: string]: { uuid: string; url: string } }> => {
  const queueName = 'demucs_requests_beta_2';
  const timeout = 30;

  const sample = await getSample({ uuid: uuid, asUser: user }, { returnLength: true });
  const parts = ['drums', 'bass', 'other', 'vocals'];
  const partCacheKeys = parts.map((part) => ({ part, cacheKey: `${sample.uuid}[${part}]` }));
  const cachedStems = await Promise.all(
    partCacheKeys.map(({ part, cacheKey }) =>
      getSample({ cacheKey, asUser: user }, { returnCompressed: true }).then((partSample) => ({
        part,
        partSample,
      }))
    )
  );
  const stemResponse = {};

  if (cachedStems.every(({ partSample }) => partSample !== null)) {
    for (const { part, partSample } of cachedStems) {
      stemResponse[part] = {
        uuid: partSample.uuid,
        url: await getSignedUrlForSampleS3Key(partSample.s3Key),
      };
    }
    return stemResponse;
  }

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

  for (const part of Object.keys(res.stems)) {
    const partS3Key = res.stems[part];
    const partCacheKey = partCacheKeys.find(({ part: p }) => p === part).cacheKey;
    const partTitleCase = part[0].toUpperCase() + part.slice(1);
    const insertedRow = await insertSample({
      name: `${partTitleCase} from ${sample.name}`,
      uuid: partS3Key.split('.')[0],
      s3Key: partS3Key,
      authorId: user.id,
      warp: sample.warp,
      cacheKey: partCacheKey,
      derivedFromSampleId: sample.id,
      derivationType: 'demucs',
      tags: [{ value: partTitleCase, category: 'Stem' }],
      lengthInMs: sample.lengthInMs,
      packId: samplePackId,
    });

    stemResponse[part] = {
      uuid: insertedRow.uuid,
      url: await getSignedUrlForSampleS3Key(partS3Key),
    };
  }

  return stemResponse;
};

const handleRequest = async (socket: Socket, message: any) => {
  const user = (socket.request as any).user;
  const sampleUUID = message.uuid;
  const onProgress = (progress) => socket.emit(SocketEvent.Progress, progress);
  return {
    stems: await runDemucs(user, sampleUUID, null, onProgress),
  };
};

export const demucsMessageHandler = async (socket: Socket, message: any) => {
  const user = (socket.request as any).user;
  await runServiceWithQuota(
    user,
    'demucs',
    async () => {
      socket.emit(SocketEvent.Response, await handleRequest(socket, message));
    },
    () => {
      socket.emit(SocketEvent.Response, { error: true, message: 'Quota exceeded' });
    }
  );
};