import query from '../server-utils/query.js';
import { declareHandler } from '../server-utils/routesHandler.js';
import { User, UserRole, isUuid } from '../types/serverTypes.js';
import { getSample } from './samples.js';
import { getUserById } from './users.js';
import { createVoiceConversionModel } from './voiceConversion.js';

const SCORE_THRESHOLD = 0.8;
/**
 * Returns latest unused script for a user,
 * returns null if no script is available
 */
const getUserLatestUnusedScript = async (
  userId: number
): Promise<{
  id: number;
  trainingScript: string;
  audioFileS3Keys: string[];
  transcribedScript: string;
  scriptTitle: string;
  matchScore: number;
} | null> => {
  const existingScript = await query(
    `
    SELECT id, training_script, transcribed_script, match_score, script_title
    FROM voice_conversion_training_scripts
    WHERE user_id = $1 AND
          model_id IS NULL AND
          status IN ('new', 'verified')
    ORDER BY created_at DESC
    `,
    [userId]
  );

  if (existingScript.rows.length && existingScript.rows[0].training_script) {
    console.log(`User ${userId} has existing unused script`);
    const scriptAudioFiles = await query(
      `
        SELECT COALESCE(s.s3_key, s.compressed_s3_key) AS sample_key
        FROM voice_conversion_training_scripts vcts
          JOIN voice_conversion_training_script_samples vctss ON vcts.id = vctss.training_script_id
          JOIN samples s ON vctss.sample_id = s.id
        WHERE vcts.id = $1
      `,
      [existingScript.rows[0].id]
    );
    const audioFileS3Keys = scriptAudioFiles.rows.map((row) => row.sample_key);
    return {
      id: existingScript.rows[0].id,
      trainingScript: existingScript.rows[0].training_script,
      audioFileS3Keys: audioFileS3Keys,
      scriptTitle: existingScript.rows[0].script_title,
      transcribedScript: existingScript.rows[0].transcribed_script,
      matchScore: existingScript.rows[0].match_score,
    };
  } else {
    console.log(`User ${userId} does not have existing unused script`);
    return null;
  }
};

const getUserScript = async (
  userId: number,
  scriptId: number
): Promise<{
  id: number;
  trainingScript: string;
  audioFileS3Keys: string[];
  scriptTitle: string;
  transcribedScript: string;
  matchScore: number;
} | null> => {
  const existingScript = await query(
    `
    SELECT id, training_script, transcribed_script, match_score
    FROM voice_conversion_training_scripts
    WHERE user_id = $1 AND
          id = $2
    ORDER BY created_at DESC
    `,
    [userId, scriptId]
  );

  if (existingScript.rows.length && existingScript.rows[0].training_script) {
    console.log(`User ${userId} has script with id ${scriptId}`);
    const scriptAudioFiles = await query(
      `
        SELECT COALESCE(s.s3_key, s.compressed_s3_key) AS sample_key
        FROM voice_conversion_training_scripts vcts
          JOIN voice_conversion_training_script_samples vctss ON vcts.id = vctss.training_script_id
          JOIN samples s ON vctss.sample_id = s.id
        WHERE vcts.id = $1
      `,
      [existingScript.rows[0].id]
    );
    const audioFileS3Keys = scriptAudioFiles.rows.map((row) => row.sample_key);
    return {
      id: existingScript.rows[0].id,
      trainingScript: existingScript.rows[0].training_script,
      audioFileS3Keys: audioFileS3Keys,
      scriptTitle: existingScript.rows[0].script_title,
      transcribedScript: existingScript.rows[0].transcribed_script,
      matchScore: existingScript.rows[0].match_score,
    };
  } else {
    console.log(`User ${userId} does not have script with id ${scriptId}`);
    return null;
  }
};

const textToWordArray = (text: string): string[] => {
  // Remove punctuation and replace with spaces
  let noPunctuation = text.replace(/[^\w\s]|_/g, ' ');

  // Replace consecutive spaces with a single space
  let singleSpaced = noPunctuation.replace(/\s+/g, ' ');

  return singleSpaced.toLowerCase().split(' ');
};

/**
 * Returns a number between 0 and 1, where 0 is no match and 1 is exact match
 */
const cosineDistance = (a: string, b: string) => {
  function getVector(words: string[]): { [key: string]: number } {
    let wordCounts = {};
    for (let word of words) {
      if (wordCounts[word]) {
        wordCounts[word]++;
      } else {
        wordCounts[word] = 1;
      }
    }
    return wordCounts;
  }

  function getDotProduct(vectorA: { [key: string]: number }, vectorB: { [key: string]: number }) {
    let dotProduct = 0;
    for (let key in vectorA) {
      if (vectorB[key]) {
        dotProduct += vectorA[key] * vectorB[key];
      }
    }
    return dotProduct;
  }

  function getMagnitude(vector: { [key: string]: number }) {
    let sumOfSquares = 0;
    for (let key in vector) {
      sumOfSquares += Math.pow(vector[key], 2);
    }
    return Math.sqrt(sumOfSquares);
  }

  const aWords = textToWordArray(a);
  const bWords = textToWordArray(b);

  const scriptVector = getVector(aWords);
  const transcriptionVector = getVector(bWords);

  const dotProduct = getDotProduct(scriptVector, transcriptionVector);
  const magnitude = getMagnitude(scriptVector) * getMagnitude(transcriptionVector);

  return dotProduct / magnitude;
};

/**
 * Returns a number between 0 and 1, where 0 is no match and 1 is exact match
 */
const levenshteinDistanceWords = (a: string, b: string): number => {
  const aWords = textToWordArray(a);
  const bWords = textToWordArray(b);
  const matrix = [];
  const sizeX = aWords.length + 1;
  const sizeY = bWords.length + 1;

  for (let x = 0; x < sizeX; x++) {
    matrix[x] = [x];
  }
  for (let y = 0; y < sizeY; y++) {
    matrix[0][y] = y;
  }

  for (let x = 1; x < sizeX; x++) {
    for (let y = 1; y < sizeY; y++) {
      let cost = aWords[x - 1] === bWords[y - 1] ? 0 : 1;
      matrix[x][y] = Math.min(matrix[x - 1][y] + 1, matrix[x][y - 1] + 1, matrix[x - 1][y - 1] + cost);
    }
  }

  return matrix[sizeX - 1][sizeY - 1];
};

const getSimilarityScore = (a: string, b: string) => {
  const cosineScore = cosineDistance(a, b);
  const distance = levenshteinDistanceWords(a, b);
  const levenshteinScore = 1 - distance / a.split(' ').length;
  return cosineScore * levenshteinScore;
};

export const TOPICS = [
  'Music',
  'Movies',
  'Art',
  'Science',
  'Architecture',
  'Astronomy',
  'Technology',
  'Engineering',
  'Mathematics',
  'Chemistry',
  'Physics',
  'Geography',
  'Transportation',
  'Food',
  'Games',
  'Hobbies',
  'Sports',
  'Animals',
  'Plants',
  'Space',
  'Languages',
];

const acceptableLanguages = {
  en: 'English',
  // ru: 'Russian',
  // ja: 'Japanese',
};

const randomTopic = () => {
  return TOPICS[Math.floor(Math.random() * TOPICS.length)];
};

const openAICall = async (prompt: string) => {
  const apiResult = await fetch('https://api.openai.com/v1/chat/completions', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      Authorization: `Bearer ${process.env.OPENAI_API_KEY}`,
    },
    body: JSON.stringify({
      model: 'gpt-3.5-turbo',
      messages: [
        {
          role: 'user',
          content: prompt,
        },
      ],
      temperature: 0.7,
    }),
  });

  const resultJson = await apiResult.json();
  const result = resultJson?.choices[0]?.message?.content;

  return result;
};

export const SAMPLES_BUCKET = 'samples-test-1';

type ClipData = {
  uuid: string;
  s3Key: string;
  buffer: Buffer;
  transcription: string;
};

const getClipDataFromSampleIds = async (sampleUUIDs: string[], user: User): Promise<ClipData[]> => {
  const clipData: ClipData[] = await Promise.all(
    sampleUUIDs.map(async (uuid) => {
      const sample = await getSample({ uuid, asUser: user }, { returnBuffer: true });
      return {
        uuid: sample.uuid,
        s3Key: sample.s3Key,
        buffer: sample.buffer,
        transcription: '',
      };
    })
  );
  return clipData;
};

const transcribeBuffer = async (buffer: Buffer): Promise<string> => {
  // Define the data for the API call
  const form = new FormData();
  form.append('file', new Blob([buffer]), 'clip.ogg');
  form.append('model', 'whisper-1');
  form.append('language', 'en');

  // Make the API call
  try {
    const result = await fetch('https://api.openai.com/v1/audio/transcriptions', {
      method: 'POST',
      body: form,
      headers: {
        Authorization: `Bearer ${process.env.OPENAI_API_KEY}`,
      },
    });

    const resultJson = await result.json();

    if (resultJson.text) {
      return resultJson.text;
    }
  } catch (error) {
    console.error(error);
  }
  return '';
};

export const generateVoiceConversionScript = declareHandler({
  func: async (req, res) => {
    // Check if the user has an existing available script
    const language = req.body.language || 'en';
    const generateNew = !!req.body.generateNew || false;
    if (!acceptableLanguages[language]) {
      return res.status(400).send({ error: 'Invalid language' });
    }

    const fullLanguage = acceptableLanguages[language];

    const unusedScript = await getUserLatestUnusedScript(req.user.id);
    if (unusedScript) {
      if (generateNew) {
        console.log('Deleting existing script');
        await query(
          `
          UPDATE voice_conversion_training_scripts
          SET status = 'deleted'
          WHERE id = $1
          `,
          [unusedScript.id]
        );
      } else {
        console.log('Returning existing script');
        return res.send({
          success: true,
          script: unusedScript.trainingScript,
          scriptId: unusedScript.id,
          scriptTitle: unusedScript.scriptTitle,
          passed: unusedScript.matchScore && unusedScript.matchScore > SCORE_THRESHOLD,
        });
      }
    }

    console.log('Generating prototype script');
    const topic = randomTopic();
    const script = await openAICall(
      `Create a 6 paragraph essay about ${topic} that can be understood by a young audience. Do not provide a title and do not provide paragraph headers. Avoid sensitive topics like religion or politics. Do not reference the reader's age. This should be written in ${fullLanguage}.`
    );
    if (!script) {
      return res.send({ success: false, passed: false, scriptId: null, message: 'Unable to generate script' });
    }

    const newScriptTitle = `A Script About ${topic}`;
    const newScript = await query(
      `
      INSERT INTO voice_conversion_training_scripts (user_id, script_title, training_script)
      VALUES ($1, $2, $3) RETURNING id
      `,
      [req.user.id, newScriptTitle, script]
    );

    if (newScript.rows.length <= 0) {
      return res.send({ success: false, passed: false, scriptId: null, message: 'Unable to generate script' });
    }

    res.send({
      success: true,
      script,
      passed: false,
      scriptId: newScript.rows[0].id,
      scriptTitle: newScriptTitle,
    });
  },
});

export const verifyVoiceConversionScript = declareHandler({
  func: async (req, res) => {
    const sampleIds: string[] = req.body.sampleIds;
    const scriptId: number = req.body.scriptId;

    const user: User = await getUserById(req.user.id);

    // return res.status(400).send({ passed: false, message: 'Verification is disabled' });
    if (!Array.isArray(sampleIds) || sampleIds.some((s) => !isUuid(s))) {
      return res.status(400).send({ error: 'Invalid request body' });
    }

    // Get the script from the latest available script
    const unusedScript = await getUserScript(req.user.id, scriptId);
    if (!unusedScript) {
      return res.send({ success: false, message: 'Unable to run verification.' });
    }

    const clipData = await getClipDataFromSampleIds(sampleIds, user);
    if (clipData.length <= 0) {
      return res.send({ success: false, message: 'Unable to run verification.' });
    }

    for (const c of clipData) {
      c.transcription = await transcribeBuffer(c.buffer);
    }

    // Sort the buffers by the order of the sampleIds (even duplicates)
    const sortedBuffers: ClipData[] = [];
    for (const sampleId of sampleIds) {
      const buffer = clipData.find((b) => b.uuid === sampleId);
      if (buffer) {
        sortedBuffers.push(buffer);
      }
    }

    // Add up all the transcriptions
    const completeTranscription = sortedBuffers.reduce((acc, buffer) => {
      return acc + buffer.transcription;
    }, '');

    if (completeTranscription.length <= 0) {
      return res.send({ success: false, message: 'Unable to run verification.' });
    }

    const score =
      user.role === UserRole.Admin ? 999 : getSimilarityScore(completeTranscription, unusedScript.trainingScript);

    console.log('score', score);

    // Delete all existing sample references
    await query(
      `
      DELETE FROM voice_conversion_training_script_samples
      WHERE training_script_id = $1
      `,
      [unusedScript.id]
    );

    console.log('deleted existing samples');

    // Insert transcription and score into database
    const status = score >= SCORE_THRESHOLD ? 'verified' : 'new';
    await query(
      `
      UPDATE voice_conversion_training_scripts
      SET transcribed_script = $1,
          match_score = $2,
          status = $3
      WHERE id = $4
      `,
      [completeTranscription, score, status, unusedScript.id]
    );

    console.log('updated script');

    // Insert sample references into database
    await query(
      `
      INSERT INTO voice_conversion_training_script_samples (training_script_id, sample_id)
      SELECT $1, id
      FROM samples
      WHERE uuid = ANY($2)
      `,
      [unusedScript.id, sampleIds]
    );

    console.log('inserted samples');

    return res.send({ success: true, passed: score >= SCORE_THRESHOLD });
  },
});

export const startTrainingModelFromScript = declareHandler({
  func: async (req, res) => {
    const { scriptId, name } = req.body;
    const scriptQuery = await query(
      `
      SELECT id
      FROM voice_conversion_training_scripts
      WHERE id = $1 AND user_id = $2 AND model_id IS NULL AND status = 'verified' AND match_score >= $3
      ORDER BY created_at DESC
      `,
      [scriptId, req.user.id, SCORE_THRESHOLD]
    );

    if (scriptQuery.rows.length <= 0) {
      return res.status(404).send({ success: false, error: 'Script not found' });
    }

    const scriptAudioFiles = await query(
      `
        SELECT s.uuid
        FROM voice_conversion_training_scripts vcts
          JOIN voice_conversion_training_script_samples vctss ON vcts.id = vctss.training_script_id
          JOIN samples s ON vctss.sample_id = s.id
        WHERE vcts.id = $1
      `,
      [scriptId]
    );

    // Get the corpus sample IDs
    const sampleIds = scriptAudioFiles.rows.map((row) => row.uuid);

    // Run the createVoiceConversionModel function
    const result = await createVoiceConversionModel(name, req.user.id, sampleIds, 'voice');

    if (result.uuid !== null && result.id !== null) {
      // Update the script with the model ID
      await query(
        `
        UPDATE voice_conversion_training_scripts
        SET model_id = $1, status = 'trained'
        WHERE id = $2
        `,
        [result.id, scriptId]
      );
      console.log('updated script with model id');
      return res.status(202).send({ success: true, uuid: result.uuid, message: 'Script training' });
    } else {
      return res.status(400).send({ success: false, error: 'Invalid Request' });
    }
  },
});