STT Renderer

'use client'
 
import TranscriptWaveGlyph from '@jamie/icons/TranscriptWaveGlyph'
import {
  getRenderableSegments,
  initialStreamState,
  reduceTranscriptionMessage,
  type StreamState
} from '@jamie/transcript-stream/reducer'
import type { SonioxToken, TranscriptionMessage } from '@jamie/transcript-stream/types'
import { Button } from '@jamie/ui/button'
import { SttRender } from '@jamie/ui/compose/stt-render'
import { Pause, Play, RotateCcw } from 'lucide-react'
import { useCallback, useEffect, useMemo, useReducer, useRef, useState } from 'react'
 
interface ChunkScript {
  text: string
  buildMs: number
  holdMs: number
}
 
interface SegmentScript {
  speaker: string
  startMs: number
  chunks: ChunkScript[]
}
 
const PROGRAM: SegmentScript[] = [
  {
    speaker: 'speaker_a',
    startMs: 0,
    chunks: [
      { text: 'Welcome to', buildMs: 350, holdMs: 350 },
      { text: " today's standup.", buildMs: 500, holdMs: 300 },
      { text: " Let's start with updates.", buildMs: 700, holdMs: 250 }
    ]
  },
  {
    speaker: 'speaker_b',
    startMs: 3500,
    chunks: [
      { text: 'I shipped the', buildMs: 400, holdMs: 250 },
      { text: ' streaming refactor yesterday', buildMs: 700, holdMs: 300 },
      { text: " and we're seeing", buildMs: 550, holdMs: 250 },
      { text: ' 40% lower latency.', buildMs: 650, holdMs: 250 }
    ]
  },
  {
    speaker: 'speaker_c',
    startMs: 9200,
    chunks: [
      { text: 'Next on my list', buildMs: 500, holdMs: 250 },
      { text: ' is the speaker diarization fix.', buildMs: 850, holdMs: 300 }
    ]
  }
]
 
const INITIAL_DELAY_MS = 1500
const CHUNK_DURATION_MS = 800
 
interface ScriptedTick {
  delayMs: number
  message: TranscriptionMessage
}
 
function makeMessage(timestamp: number, tokens: SonioxToken[]): TranscriptionMessage {
  return {
    type: 'transcription',
    timestamp,
    data: {
      tokens,
      final_audio_proc_ms: timestamp,
      total_audio_proc_ms: timestamp
    }
  }
}
 
function buildTimeline() {
  const ticks: ScriptedTick[] = []
  let cursor = INITIAL_DELAY_MS
 
  for (let segIdx = 0; segIdx < PROGRAM.length; segIdx++) {
    const seg = PROGRAM[segIdx]
    if (!seg) continue
    let chunkStartMs = seg.startMs
 
    for (let chunkIdx = 0; chunkIdx < seg.chunks.length; chunkIdx++) {
      const chunk = seg.chunks[chunkIdx]
      if (!chunk) continue
      const chunkEndMs = chunkStartMs + CHUNK_DURATION_MS
      const needsBoundary = chunkIdx === 0 && segIdx > 0
 
      const buildTokens: SonioxToken[] = []
      // To start a new segment, the reducer needs a final token with a new
      // speaker. Send an empty boundary token alongside the first chunk's
      // non-final preview so the previous segment finalizes cleanly while the
      // new segment immediately enters its build phase.
      if (needsBoundary) {
        buildTokens.push({
          text: '',
          start_ms: seg.startMs,
          end_ms: seg.startMs,
          confidence: 1,
          is_final: true,
          speaker: seg.speaker
        })
      }
      buildTokens.push({
        text: chunk.text,
        start_ms: chunkStartMs,
        end_ms: chunkEndMs,
        confidence: 0.85,
        is_final: false,
        speaker: seg.speaker
      })
 
      ticks.push({ delayMs: cursor, message: makeMessage(cursor, buildTokens) })
      cursor += chunk.buildMs
 
      ticks.push({
        delayMs: cursor,
        message: makeMessage(cursor, [
          {
            text: chunk.text,
            start_ms: chunkStartMs,
            end_ms: chunkEndMs,
            confidence: 0.97,
            is_final: true,
            speaker: seg.speaker
          }
        ])
      })
      cursor += chunk.holdMs
 
      chunkStartMs = chunkEndMs
    }
  }
 
  ticks.push({ delayMs: cursor, message: { type: 'stt_done', timestamp: cursor } })
  return ticks
}
 
const TIMELINE = buildTimeline()
 
type Action = { kind: 'message'; message: TranscriptionMessage } | { kind: 'reset' }
 
function streamReducer(state: StreamState, action: Action) {
  if (action.kind === 'reset') return initialStreamState
  return reduceTranscriptionMessage(state, action.message)
}
 
export function SttRendererDemo() {
  const [state, dispatch] = useReducer(streamReducer, initialStreamState)
  const [isPlaying, setIsPlaying] = useState(true)
  const tickIndex = useRef(0)
  const timeoutId = useRef<ReturnType<typeof setTimeout> | null>(null)
 
  const stop = useCallback(() => {
    if (timeoutId.current !== null) {
      clearTimeout(timeoutId.current)
      timeoutId.current = null
    }
  }, [])
 
  const advance = useCallback(() => {
    const tick = TIMELINE[tickIndex.current]
    if (!tick) {
      setIsPlaying(false)
      return
    }
    dispatch({ kind: 'message', message: tick.message })
    tickIndex.current += 1
    const next = TIMELINE[tickIndex.current]
    if (!next) {
      setIsPlaying(false)
      return
    }
    timeoutId.current = setTimeout(advance, next.delayMs - tick.delayMs)
  }, [])
 
  const handleReset = useCallback(() => {
    stop()
    tickIndex.current = 0
    dispatch({ kind: 'reset' })
    setIsPlaying(true)
    timeoutId.current = setTimeout(advance, INITIAL_DELAY_MS)
  }, [advance, stop])
 
  const handleToggle = useCallback(() => {
    if (isPlaying) {
      stop()
      setIsPlaying(false)
      return
    }
    if (tickIndex.current >= TIMELINE.length) {
      tickIndex.current = 0
      dispatch({ kind: 'reset' })
      setIsPlaying(true)
      timeoutId.current = setTimeout(advance, INITIAL_DELAY_MS)
      return
    }
    setIsPlaying(true)
    advance()
  }, [advance, isPlaying, stop])
 
  useEffect(() => {
    timeoutId.current = setTimeout(advance, INITIAL_DELAY_MS)
    return () => {
      stop()
      tickIndex.current = 0
    }
  }, [advance, stop])
 
  const { segments, interim } = useMemo(() => getRenderableSegments(state), [state])
  const isEmpty = segments.length === 0 && !interim
 
  return (
    <div className="relative w-full max-w-md min-h-72 rounded-md border border-outline bg-surface p-4">
      <SttRender.List>
        <SttRender.Placeholder show={isEmpty} />
        {segments.map((seg) => (
          <SttRender.Segment key={seg.id}>
            <SttRender.Timestamp startMs={seg.startMs} endMs={seg.endMs} />
            <SttRender.Text>{seg.text}</SttRender.Text>
          </SttRender.Segment>
        ))}
        {interim && (
          <SttRender.Segment>
            <SttRender.Timestamp startMs={interim.startMs} endMs={interim.endMs} />
            <SttRender.InterimText
              finalText={interim.finalText}
              nonFinalText={interim.nonFinalText}
            >
              <TranscriptWaveGlyph className="ml-1 inline-block align-middle text-on-surface" />
            </SttRender.InterimText>
          </SttRender.Segment>
        )}
      </SttRender.List>
      <div className="absolute bottom-2 right-2 flex">
        <Button variant="ghost" size="icon-sm" onClick={handleReset} aria-label="Reset">
          <RotateCcw className="size-4" />
        </Button>
        <Button
          variant="ghost"
          size="icon-sm"
          onClick={handleToggle}
          aria-label={isPlaying ? 'Pause' : 'Play'}
        >
          {isPlaying ? <Pause className="size-4" /> : <Play className="size-4" />}
        </Button>
      </div>
    </div>
  )
}