CedarCopilot/cedar-mail

import { useState, useRef, useCallback, useEffect } from 'react';
import { useTRPC } from '@/providers/query-provider';
import { useMutation } from '@tanstack/react-query';
import { toast } from 'sonner';

export function useGroqDictation(onTranscript: (text: string) => void, options?: { format?: boolean }) {
  const [isRecording, setIsRecording] = useState(false);
  const [isProcessing, setIsProcessing] = useState(false);
  const [isFormatting, setIsFormatting] = useState(false);

  const recorderRef = useRef<MediaRecorder | null>(null);
  const chunksRef = useRef<Blob[]>([]);
  const streamRef = useRef<MediaStream | null>(null);
  const onTranscriptRef = useRef(onTranscript);
  const formatRef = useRef(options?.format ?? false);

  useEffect(() => {
    onTranscriptRef.current = onTranscript;
  }, [onTranscript]);

  const trpc = useTRPC();
  const { mutateAsync: transcribeAudio } = useMutation(trpc.transcription.transcribeAudio.mutationOptions());
  const { mutateAsync: formatTranscription } = useMutation(trpc.transcription.formatTranscription.mutationOptions());
  const transcribeAudioRef = useRef(transcribeAudio);
  const formatTranscriptionRef = useRef(formatTranscription);
  useEffect(() => {
    transcribeAudioRef.current = transcribeAudio;
  }, [transcribeAudio]);
  useEffect(() => {
    formatTranscriptionRef.current = formatTranscription;
  }, [formatTranscription]);

  const start = useCallback(async () => {
    try {
      const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
      streamRef.current = stream;

      const mimeType = MediaRecorder.isTypeSupported('audio/webm;codecs=opus')
        ? 'audio/webm;codecs=opus'
        : MediaRecorder.isTypeSupported('audio/webm')
          ? 'audio/webm'
          : 'audio/ogg';

      const recorder = new MediaRecorder(stream, { mimeType });
      recorderRef.current = recorder;
      chunksRef.current = [];

      recorder.ondataavailable = (e) => {
        if (e.data.size > 0) chunksRef.current.push(e.data);
      };

      recorder.onstop = async () => {
        stream.getTracks().forEach((t) => t.stop());
        streamRef.current = null;

        const blob = new Blob(chunksRef.current, { type: mimeType });
        if (blob.size < 1000) {
          setIsRecording(false);
          return;
        }

        setIsProcessing(true);
        setIsRecording(false);

        try {
          const arrayBuffer = await blob.arrayBuffer();
          const audio = btoa(
            new Uint8Array(arrayBuffer).reduce(
              (data, byte) => data + String.fromCharCode(byte),
              '',
            ),
          );

          const data = await transcribeAudioRef.current({ audio, mimeType });
          if (data.text) {
            if (formatRef.current) {
              setIsProcessing(false);
              setIsFormatting(true);
              const formatted = await formatTranscriptionRef.current({ text: data.text });
              console.log('[dictation] raw:', data.text);
              console.log('[dictation] formatted:', formatted.html);
              onTranscriptRef.current(formatted.html);
            } else {
              onTranscriptRef.current(data.text);
            }
          }
        } catch (err) {
          toast.error(err instanceof Error ? err.message : 'Transcription failed');
        } finally {
          setIsProcessing(false);
          setIsFormatting(false);
        }
      };

      recorder.start(250);
      setIsRecording(true);
    } catch (err) {
      toast.error(err instanceof Error ? err.message : 'Could not access microphone');
    }
  }, []);

  const stop = useCallback(() => {
    if (recorderRef.current && recorderRef.current.state !== 'inactive') {
      recorderRef.current.stop();
    }
  }, []);

  const toggle = useCallback(() => {
    if (isRecording) stop();
    else void start();
  }, [isRecording, start, stop]);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      if (recorderRef.current && recorderRef.current.state !== 'inactive') {
        recorderRef.current.stop();
      }
      streamRef.current?.getTracks().forEach((t) => t.stop());
    };
  }, []);

  return { isRecording, isProcessing, isFormatting, toggle };
}