import { Socket } from 'socket.io';
import { v4 as uuidv4 } from 'uuid';
import { getRedisClient, withIsolatedRedisClient } from '../external-services/redis.js';
import query, { transaction } from '../server-utils/query.js';
import { declareHandler } from '../server-utils/routesHandler.js';
import { SocketEvent } from '../server-utils/socketMessageHandler.js';
import { isUuid } from '../types/serverTypes.js';
import { hashJson } from '../utils/hash-json.js';
import runServiceWithQuota from '../utils/runServiceWithQuota.js';
import { runDemucs, subtract } from './demucs.js';
import { getSample, getSignedUrlForSampleS3Key, insertSample } from './samples.js';
import { getUserById } from './users.js';

/**
 * This file contains route handlers and websocket message handlers for
 * the voice conversion service. The ML model implementation is in ml/rvc
 * and communicates with this service via Redis (for inference) + Postgres
 * (for training).
 *
 */

const inferenceQueueName = 'voice_conversion_requests_2';

const localTrainingEnabled = !['production', 'staging', 'beta'].includes(process.env.ENVIRONMENT);

export const getVoiceConversionModels = declareHandler({
  func: async (req, res) => {
    const models = await query(
      `
      SELECT vcm.*, r.*, coalesce(corpus_samples.sample_uuids, '{}') corpus_samples
       FROM voice_conversion_models vcm
       LEFT JOIN LATERAL (
        SELECT * FROM voice_conversion_model_training_runs r
        WHERE r.model_id = vcm.id
        ORDER BY r.created_at DESC LIMIT 1
       ) r ON true
       LEFT JOIN LATERAL (
        SELECT array_agg(s.uuid) sample_uuids FROM voice_conversion_model_corpus c
         JOIN samples s ON s.id = c.sample_id
         WHERE c.model_id = vcm.id
       ) corpus_samples ON true
       WHERE (author_id = $1 OR author_id IS NULL) AND NOT archived
      `,
      [req.user.id]
    );

    if (localTrainingEnabled) {
      const client = getRedisClient();
      for (let model of models.rows) {
        if (!['pending', 'training'].includes(model.state)) {
          continue;
        }
        const queueName = `rvc_progress:${model.uuid}`;
        while (true) {
          const message = await client.lpop(queueName);
          if (!message) {
            break;
          }
          const { progress, state, s3Key, error } = JSON.parse(message);
          await query(
            `
            UPDATE voice_conversion_model_training_runs
            SET (progress, state, s3_key, error) = (coalesce($2, progress), coalesce($3, state), coalesce($4, s3_key), coalesce($5, error))
            WHERE id = $1
            `,
            [model.id, progress, state, s3Key, error]
          );
        }
      }
    }

    return res.send(
      models.rows.map((model) => ({
        uuid: model.uuid,
        name: model.name,
        createdAt: model.created_at,
        state: model.state,
        progress: model.progress,
        corpusSamples: model.corpus_samples,
      }))
    );
  },
});

export const createVoiceConversionModel = async (
  name: string,
  userId: number,
  corpusSamples: string[],
  modelType: 'instrument' | 'voice' = 'instrument'
): Promise<{ error: string | null; uuid: string | null; id: number | null }> => {
  if (
    typeof name !== 'string' ||
    !Array.isArray(corpusSamples) ||
    corpusSamples.some((s) => !isUuid(s)) ||
    corpusSamples.length === 0
  ) {
    return { error: 'invalid', uuid: null, id: null };
  }

  const user = await getUserById(userId);

  const deduplicatedCorpusSamples = [...new Set(corpusSamples.map((x) => x.toLowerCase()))] as string[];

  // voice only, or all-but-voice
  const stemSeparatedCorpusSamples = await Promise.all(
    deduplicatedCorpusSamples.map(async (s) => {
      const stems = await runDemucs(user, s);
      if (modelType === 'voice') {
        return (stems as any).vocals.uuid;
      } else {
        return await subtract(s, (stems as any).vocals.uuid, user);
      }
    })
  );

  const corpusSampleIds = await Promise.all(
    stemSeparatedCorpusSamples.map(async (uuid) => {
      const sample = await getSample({ uuid, asUser: user }, {});
      return `${sample.id}`;
    })
  );

  if (corpusSampleIds.length !== stemSeparatedCorpusSamples.length) {
    return { error: 'notfound', uuid: null, id: null };
  }

  const transactionResult = await transaction(async (query) => {
    const modelInsertQuery = await query(
      'INSERT INTO voice_conversion_models (author_id, name, model_type) VALUES ($1, $2, $3) ON CONFLICT (author_id, name) WHERE author_id IS NOT NULL AND NOT archived DO NOTHING RETURNING id, uuid',
      [userId, name, modelType]
    );

    if (modelInsertQuery.rows.length === 0) {
      return { error: 'exists', uuid: null, id: null };
    }

    const modelId = modelInsertQuery.rows[0].id as number;

    await query(
      'INSERT INTO voice_conversion_model_corpus (model_id, sample_id) SELECT $1, u FROM unnest($2::integer[]) u',
      [modelId, corpusSampleIds]
    );

    await enqueueModelTrainingRequest(modelId, query);

    return { error: null, uuid: modelInsertQuery.rows[0].uuid as string, id: modelId };
  });

  return transactionResult;
};

export const createVoiceConversionModelEndpoint = declareHandler({
  func: async (req, res) => {
    const { name, corpusSamples } = req.body;
    const result = await createVoiceConversionModel(name, req.user.id, corpusSamples);

    if (result.uuid !== null) {
      return res.send({ uuid: result.uuid });
    } else if (result.error !== null) {
      switch (result.error) {
        case 'invalid':
          return res.status(400).send({ error: 'Invalid request body' });
        case 'notfound':
          return res.status(404).send({ error: 'Corpus sample not found' });
        case 'exists':
          return res.status(400).send({ error: 'Model with this name already exists' });
        default:
          return res.status(500).send({ error: 'Unknown error' });
      }
    } else {
      return res.status(500).send({ error: 'Unknown error' });
    }
  },
});

export const deleteVoiceConversionModel = declareHandler({
  func: async (req, res) => {
    const { uuid } = req.params;

    if (!isUuid(uuid)) {
      return res.status(400).send({ error: 'Invalid request params' });
    }

    await transaction(async (query) => {
      const archiveQuery = await query(
        'UPDATE voice_conversion_models SET archived = true WHERE uuid = $1 AND author_id = $2 RETURNING id',
        [uuid, req.user.id]
      );

      if (archiveQuery.rows.length === 0) {
        return res.status(404).send({ error: 'Model not found' });
      }

      await query(
        `UPDATE voice_conversion_model_training_runs SET state = 'cancelled' FROM voice_conversion_models vcm WHERE model_id = vcm.id AND vcm.uuid = $1 AND state IN ('pending', 'training')`,
        [uuid]
      );

      if (localTrainingEnabled) {
        const client = getRedisClient();
        await client.lpush(`rvc_cancel:${uuid}`, '1');
      }
    }, 5);

    return res.status(204).send();
  },
});

const promiseSpawn = async (command: string, args: string[], data: string): Promise<void> => {
  const { spawn } = await import('child_process');
  return new Promise<void>((resolve, reject) => {
    const proc = spawn(command, args);
    proc.stdin.write(data);
    proc.stdin.end();
    proc.stdout.on('data', (data) => {
      console.log(`stdout: ${data}`);
    });
    proc.stderr.on('data', (data) => {
      console.error(`stderr: ${data}`);
    });
    proc.on('error', reject);
    proc.on('exit', (code) => {
      if (code === 0) {
        resolve();
      } else {
        reject(new Error(`Process exited with code ${code}`));
      }
    });
  });
};

const enqueueModelTrainingRequestLocal = async (modelId: number, txnQuery: typeof query = query) => {
  const modelQuery = await txnQuery(
    `
    SELECT m.uuid, array_agg(distinct coalesce(s.s3_key, s.compressed_s3_key)) sample_s3_keys FROM voice_conversion_models m JOIN voice_conversion_model_corpus c ON c.model_id = m.id JOIN samples s ON s.id = c.sample_id WHERE m.id = $1 GROUP BY m.id, m.uuid
    `,
    [modelId]
  );
  const { uuid: modelUuid, sample_s3_keys: corpusS3Keys } = modelQuery.rows[0];

  // big hack alert - shell out to kubectl to run the training job

  await promiseSpawn(
    'kubectl',
    ['create', '-f', '-'],
    JSON.stringify({
      apiVersion: 'batch/v1',
      kind: 'Job',
      metadata: {
        name: `rvc-trainer-local-${modelUuid}`,
        namespace: 'production',
      },
      spec: {
        parallelism: 1,
        completions: 1,
        activeDeadlineSeconds: 3600,
        backoffLimit: 3,
        template: {
          spec: {
            restartPolicy: 'Never',
            nodeSelector: {
              'karpenter.sh/capacity-type': 'on-demand',
              'karpenter.k8s.aws/instance-family': 'g5',
            },
            serviceAccountName: 'composer-service-account',
            hostIPC: true,
            containers: [
              {
                name: 'rvc',
                image: '035363416972.dkr.ecr.us-east-1.amazonaws.com/minimal:latest',
                command: ['/opt/entrypoint.sh', '--local-train', corpusS3Keys.join(','), modelUuid],
                env: [
                  {
                    name: 'RVC_MODELS_S3_BUCKET',
                    value: 'wavtool-user-rvc-models',
                  },
                  {
                    name: 'RVC_SAMPLES_S3_BUCKET',
                    value: 'samples-test-1',
                  },
                  {
                    name: 'RVC_CACHE_DIR',
                    value: '/tmp/rvc_cache',
                  },
                  { name: 'QUEUE_NAME', value: 'whatever' },
                  { name: 'REDIS_HOST', value: 'prod-redis.redis.svc.cluster.local' },
                  { name: 'REDIS_PORT', value: '26379' },
                  { name: 'REDIS_SENTINEL_NAME', value: 'wavtoolmaster' },
                ],
                envFrom: [
                  {
                    secretRef: {
                      name: 'redis-secret',
                    },
                  },
                ],
                volumeMounts: [
                  {
                    name: 'persistent-storage',
                    mountPath: '/opt',
                    subPath: 'rvc_production_2023-11-21',
                  },
                  {
                    name: 'memory-backed-temp',
                    mountPath: '/tmp',
                  },
                ],
                resources: {
                  limits: {
                    'nvidia.com/gpu': 1,
                  },
                },
              },
            ],
            volumes: [
              {
                name: 'persistent-storage',
                persistentVolumeClaim: {
                  claimName: 'efs-claim',
                },
              },
              {
                name: 'memory-backed-temp',
                emptyDir: {
                  medium: 'Memory',
                  sizeLimit: '4Gi',
                },
              },
            ],
          },
        },
      },
    })
  );
};

const enqueueModelTrainingRequest = async (modelId: number, txnQuery: typeof query = query) => {
  await txnQuery(`INSERT INTO voice_conversion_model_training_runs (model_id) VALUES ($1)`, [modelId]);

  if (localTrainingEnabled) {
    await enqueueModelTrainingRequestLocal(modelId, txnQuery);
  }
};

const handleRequest = async (socket: Socket, message: any) => {
  const timeout = 60;
  const user = (socket.request as any).user;

  if (!user.features.ExecutePrototype) {
    return { error: true, message: 'Please upgrade to use this feature' };
  }

  const sample = await getSample({ uuid: message.sampleUuid, asUser: user }, { returnLength: true });

  const modelQuery = await query(
    `
    SELECT vcm.id, r.s3_key, vcm.name
     FROM voice_conversion_models vcm
     LEFT JOIN voice_conversion_model_training_runs r ON r.model_id = vcm.id
     WHERE
      vcm.uuid = $1::uuid
      AND (vcm.author_id = $2 OR vcm.author_id IS NULL)
      AND NOT vcm.archived
      AND r.state = 'finished'
     ORDER BY r.created_at DESC LIMIT 1
    `,
    [message.modelUuid, user.id]
  );

  if (modelQuery.rows.length === 0) {
    return { error: true, message: 'Model not found' };
  }

  const { s3_key: modelS3Key, name: modelName, id: modelId } = modelQuery.rows[0];

  if (modelS3Key === null) {
    return { error: true, message: 'Model not trained' };
  }

  const requestData = {
    sampleS3Key: sample.s3Key,
    sampleDuration: sample.lengthInMs / 1000.0 || null,
    modelS3Key,
  } as any;
  if (message.f0Delta !== undefined) {
    requestData.f0Delta = message.f0Delta;
  }
  if (message.pitchCurve !== undefined) {
    if (requestData.f0Delta !== undefined) {
      return { error: true, message: 'Cannot set both f0Delta and pitchCurve' };
    }
    requestData.pitchCurve = message.pitchCurve;
  }

  const cacheKey = hashJson(requestData);
  const cachedSample = await getSample({ cacheKey, asUser: user }, { returnCompressed: true });

  if (cachedSample !== null) {
    return {
      uuid: cachedSample.uuid,
      url: await getSignedUrlForSampleS3Key(cachedSample.s3Key),
      pitchCurve: cachedSample.analysis.pitchCurve,
    };
  }

  const redisClient = getRedisClient();
  const now = Date.now() / 1000.0;
  const requestUuid = uuidv4();

  await redisClient.lpush(
    inferenceQueueName,
    JSON.stringify({
      id: requestUuid,
      data: requestData,
      timeSubmitted: now,
    })
  );
  const res = await withIsolatedRedisClient(timeout * 1000 + 2000, async (isolatedClientPromise) => {
    const isolatedClient = await isolatedClientPromise;
    while (true) {
      const res = await isolatedClient.blpop(requestUuid, timeout);
      if (!res) {
        throw new Error('Timeout waiting for response');
      }
      const message = JSON.parse(res[1]);
      if (message.progress !== undefined) {
        socket.emit(SocketEvent.Progress, message.progress);
      } else {
        return message;
      }
    }
  });
  if (res.error) {
    throw new Error(res.error);
  }

  const { s3Key, pitchCurve } = res;

  const insertedRow = await insertSample({
    name: `${sample.name} as ${modelName}`,
    uuid: s3Key.split('.')[0],
    s3Key,
    authorId: user.id,
    cacheKey,
    analysis: { pitchCurve },
    derivedFromSampleId: sample.id,
    derivationType: 'voiceConversion',
    tags: [
      {
        category: 'Type',
        value: 'Voice Conversion Output',
      },
    ],
  });

  await query(
    `
    INSERT INTO sample_voice_conversion_models (sample_id, voice_conversion_model_id) VALUES ($1, $2);`,
    [insertedRow.id, modelId]
  );

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

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