front/src/components/PodcastGenerator.tsx

import { useEffect, useState } from 'react';
import { AudioPlayer } from './AudioPlayer';
import { Podcast, PodcastTurn } from '../utils/types';
import { parse } from 'yaml';
import {
  generateAudio,
  joinAudio,
  loadWavAndDecode,
  pickRand,
} from '../utils/utils';

interface GenerationStep {
  turn: PodcastTurn;
  audioBuffer?: AudioBuffer;
}

const SPEEDS = [
  { name: 'slow AF', value: 0.8 },
  { name: 'slow', value: 0.9 },
  { name: 'a bit slow', value: 1.0 },
  { name: 'most natural', value: 1.1 },
  { name: 'a bit fast', value: 1.2 },
  { name: 'fast!', value: 1.3 },
  { name: 'fast AF', value: 1.4 },
];

const SPEAKERS = [
  { name: '🇺🇸 🚺 Heart ❤️', value: 'af_heart' },
  { name: '🇺🇸 🚺 Bella 🔥', value: 'af_bella' },
  { name: '🇺🇸 🚺 Nicole 🎧', value: 'af_nicole' },
  { name: '🇺🇸 🚺 Aoede', value: 'af_aoede' },
  { name: '🇺🇸 🚺 Kore', value: 'af_kore' },
  { name: '🇺🇸 🚺 Sarah', value: 'af_sarah' },
  { name: '🇺🇸 🚺 Nova', value: 'af_nova' },
  { name: '🇺🇸 🚺 Sky', value: 'af_sky' },
  { name: '🇺🇸 🚺 Alloy', value: 'af_alloy' },
  { name: '🇺🇸 🚺 Jessica', value: 'af_jessica' },
  { name: '🇺🇸 🚺 River', value: 'af_river' },
  { name: '🇺🇸 🚹 Michael', value: 'am_michael' },
  { name: '🇺🇸 🚹 Fenrir', value: 'am_fenrir' },
  { name: '🇺🇸 🚹 Puck', value: 'am_puck' },
  { name: '🇺🇸 🚹 Echo', value: 'am_echo' },
  { name: '🇺🇸 🚹 Eric', value: 'am_eric' },
  { name: '🇺🇸 🚹 Liam', value: 'am_liam' },
  { name: '🇺🇸 🚹 Onyx', value: 'am_onyx' },
  { name: '🇺🇸 🚹 Santa', value: 'am_santa' },
  { name: '🇺🇸 🚹 Adam', value: 'am_adam' },
  { name: '🇬🇧 🚺 Emma', value: 'bf_emma' },
  { name: '🇬🇧 🚺 Isabella', value: 'bf_isabella' },
  { name: '🇬🇧 🚺 Alice', value: 'bf_alice' },
  { name: '🇬🇧 🚺 Lily', value: 'bf_lily' },
  { name: '🇬🇧 🚹 George', value: 'bm_george' },
  { name: '🇬🇧 🚹 Fable', value: 'bm_fable' },
  { name: '🇬🇧 🚹 Lewis', value: 'bm_lewis' },
  { name: '🇬🇧 🚹 Daniel', value: 'bm_daniel' },
];

const getRandomSpeakerPair = (): { s1: string; s2: string } => {
  const s1Gender = Math.random() > 0.5 ? '🚺' : '🚹';
  const s2Gender = s1Gender === '🚺' ? '🚹' : '🚺';
  const s1 = pickRand(
    SPEAKERS.filter((s) => s.name.includes(s1Gender) && s.name.includes('🇺🇸'))
  ).value;
  const s2 = pickRand(
    SPEAKERS.filter((s) => s.name.includes(s2Gender) && s.name.includes('🇺🇸'))
  ).value;
  return { s1, s2 };
};

export const PodcastGenerator = ({
  genratedScript,
  setBusy,
  busy,
}: {
  genratedScript: string;
  setBusy: (busy: boolean) => void;
  busy: boolean;
}) => {
  const [wav, setWav] = useState<AudioBuffer | null>(null);
  const [numSteps, setNumSteps] = useState<number>(0);
  const [numStepsDone, setNumStepsDone] = useState<number>(0);

  const [script, setScript] = useState<string>('');
  const [speaker1, setSpeaker1] = useState<string>('');
  const [speaker2, setSpeaker2] = useState<string>('');
  const [speed, setSpeed] = useState<string>('1.1');

  const setRandSpeaker = () => {
    const { s1, s2 } = getRandomSpeakerPair();
    setSpeaker1(s1);
    setSpeaker2(s2);
  };
  useEffect(setRandSpeaker, []);

  useEffect(() => {
    setScript(genratedScript);
  }, [genratedScript]);

  const generatePodcast = async () => {
    setWav(null);
    setBusy(true);
    try {
      const podcast: Podcast = parse(script);
      const { speakerNames, turns } = podcast;
      for (const turn of turns) {
        // normalize it
        turn.nextGapMilisecs =
          Math.max(-600, Math.min(300, turn.nextGapMilisecs)) - 100;
        turn.text = turn.text
          .trim()
          .replace(/’/g, "'")
          .replace(/“/g, '"')
          .replace(/”/g, '"');
      }
      const steps: GenerationStep[] = turns.map((turn) => ({ turn }));
      setNumSteps(steps.length);
      setNumStepsDone(0);
      let outputWav: AudioBuffer;
      for (let i = 0; i < steps.length; i++) {
        const step = steps[i];
        const speakerIdx = speakerNames.indexOf(
          step.turn.speakerName as string
        ) as 1 | 0;
        const speakerVoice = speakerIdx === 0 ? speaker1 : speaker2;
        const url = await generateAudio(
          step.turn.text,
          speakerVoice,
          parseFloat(speed)
        );
        step.audioBuffer = await loadWavAndDecode(url);
        if (i === 0) {
          outputWav = step.audioBuffer;
        } else {
          const lastStep = steps[i - 1];
          outputWav = joinAudio(
            outputWav!,
            step.audioBuffer,
            lastStep.turn.nextGapMilisecs / 1000
          );
        }
        setNumStepsDone(i + 1);
      }
      setWav(outputWav! ?? null);
    } catch (e) {
      console.error(e);
      alert(`Error: ${(e as any).message}`);
      setWav(null);
    }
    setBusy(false);
    setNumStepsDone(0);
    setNumSteps(0);
  };

  const isGenerating = numSteps > 0;

  return (
    <>
      <div className="card bg-base-100 w-full shadow-xl">
        <div className="card-body">
          <h2 className="card-title">Step 2: Script (YAML format)</h2>
          <textarea
            className="textarea textarea-bordered w-full h-72 p-2"
            placeholder="Type your script here..."
            value={script}
            onChange={(e) => setScript(e.target.value)}
          ></textarea>

          <div className="grid grid-cols-2 gap-4">
            <label className="form-control w-full">
              <div className="label">
                <span className="label-text">Speaker 1</span>
              </div>
              <select
                className="select select-bordered"
                value={speaker1}
                onChange={(e) => setSpeaker1(e.target.value)}
              >
                {SPEAKERS.map((s) => (
                  <option key={s.value} value={s.value}>
                    {s.name}
                  </option>
                ))}
              </select>
            </label>
            <label className="form-control w-full">
              <div className="label">
                <span className="label-text">Speaker 2</span>
              </div>
              <select
                className="select select-bordered"
                value={speaker2}
                onChange={(e) => setSpeaker2(e.target.value)}
              >
                {SPEAKERS.map((s) => (
                  <option key={s.value} value={s.value}>
                    {s.name}
                  </option>
                ))}
              </select>
            </label>
          </div>

          <div className="grid grid-cols-2 gap-4">
            <button className="btn" onClick={setRandSpeaker}>
              Randomize speakers
            </button>
            <label className="form-control w-full">
              <select
                className="select select-bordered"
                value={speed.toString()}
                onChange={(e) => setSpeed(e.target.value)}
              >
                {SPEEDS.map((s) => (
                  <option key={s.value} value={s.value.toString()}>
                    Speed: {s.name} ({s.value})
                  </option>
                ))}
              </select>
            </label>
          </div>

          <button
            className="btn btn-primary mt-2"
            onClick={generatePodcast}
            disabled={busy || !script || isGenerating}
          >
            {isGenerating ? (
              <>
                <span className="loading loading-spinner loading-sm"></span>
                Generating ({numStepsDone}/{numSteps})...
              </>
            ) : (
              'Generate podcast'
            )}
          </button>

          {isGenerating && (
            <progress
              className="progress progress-primary mt-2"
              value={numStepsDone}
              max={numSteps}
            ></progress>
          )}
        </div>
      </div>

      {wav && (
        <div className="card bg-base-100 w-full shadow-xl">
          <div className="card-body">
            <h2 className="card-title">Step 3: Listen to your podcast</h2>
            <AudioPlayer audioBuffer={wav} />
          </div>
        </div>
      )}
    </>
  );
};