cinny/src/app/features/settings/general/DenoiseTester.tsx

import React, { useCallback, useEffect, useRef, useState } from 'react';
import { Box, Button, color, config, Text } from 'folds';
import { DenoiseModelId } from '../../../state/settings';
import { DENOISE_MODELS } from '../../../utils/lotusDenoiseUtils';
import {
  DenoiseNode,
  buildGateNode,
  buildModelNode,
  readDb,
  sampleRateFor,
} from '../../../utils/denoisePipeline';

const MAX_RECORD_MS = 6000;

// Live monitor mirrors the call's capture (respects the user's native-NS choice).
const MIC_CONSTRAINTS = (nativeNS: boolean): MediaStreamConstraints => ({
  audio: {
    noiseSuppression: nativeNS,
    channelCount: 1,
    echoCancellation: true,
    autoGainControl: true,
  },
});

// Record & compare captures fully RAW audio (no browser noise suppression / AGC
// / echo cancel) so each model's effect on real background noise is audible.
// Capturing with native NS on would pre-clean the clip and make Raw/RNNoise/
// Speex sound identical.
const RAW_CONSTRAINTS: MediaStreamConstraints = {
  audio: {
    noiseSuppression: false,
    echoCancellation: false,
    autoGainControl: false,
    channelCount: 1,
  },
};

/** A -100..0 dBFS bar with optional threshold marker. */
function DbMeter({ label, db, threshold }: { label: string; db: number; threshold?: number }) {
  const pct = Math.max(0, Math.min(100, db + 100));
  const markerPct = threshold === undefined ? null : Math.max(0, Math.min(100, threshold + 100));
  return (
    <Box direction="Column" gap="100">
      <Box direction="Row" justifyContent="SpaceBetween">
        <Text size="T200">{label}</Text>
        <Text size="T200">{db <= -100 ? '−∞ dB' : `${Math.round(db)} dB`}</Text>
      </Box>
      <Box
        style={{
          position: 'relative',
          height: '12px',
          background: color.Surface.Container,
          border: `${config.borderWidth.B300} solid ${color.Surface.ContainerLine}`,
          borderRadius: '6px',
          overflow: 'hidden',
        }}
      >
        <Box
          style={{
            position: 'absolute',
            top: 0,
            left: 0,
            bottom: 0,
            width: `${pct}%`,
            background: color.Success.Main,
            transition: 'width 0.05s linear',
          }}
        />
        {markerPct !== null && (
          <Box
            style={{
              position: 'absolute',
              top: 0,
              bottom: 0,
              left: `${markerPct}%`,
              width: '2px',
              background: color.Primary.Main,
            }}
          />
        )}
      </Box>
    </Box>
  );
}

type DenoiseTesterProps = {
  model: DenoiseModelId;
  useGate: boolean;
  gateThreshold: number;
  nativeNS: boolean;
};

export function DenoiseTester({ model, useGate, gateThreshold, nativeNS }: DenoiseTesterProps) {
  // ── Live loopback monitor ────────────────────────────────────────────────
  const liveRef = useRef<{
    ctx: AudioContext;
    stream: MediaStream;
    model: DenoiseNode;
    gate: AudioWorkletNode | null;
    raf: number;
    inAnalyser: AnalyserNode;
    outAnalyser: AnalyserNode;
  } | null>(null);
  const [live, setLive] = useState(false);
  const [inDb, setInDb] = useState(-100);
  const [outDb, setOutDb] = useState(-100);

  const stopLive = useCallback(() => {
    const s = liveRef.current;
    liveRef.current = null;
    if (s) {
      cancelAnimationFrame(s.raf);
      try {
        s.gate?.disconnect();
        s.model.dispose();
        s.model.node.disconnect();
      } catch {
        /* noop */
      }
      s.stream.getTracks().forEach((t) => t.stop());
      s.ctx.close().catch(() => undefined);
    }
    setLive(false);
    setInDb(-100);
    setOutDb(-100);
  }, []);

  const startLive = async () => {
    try {
      const ctx = new AudioContext({ sampleRate: sampleRateFor(model) });
      const stream = await navigator.mediaDevices.getUserMedia(MIC_CONSTRAINTS(nativeNS));
      const source = ctx.createMediaStreamSource(stream);
      const inAnalyser = ctx.createAnalyser();
      inAnalyser.fftSize = 1024;
      source.connect(inAnalyser);

      let head: AudioNode = source;
      let gate: AudioWorkletNode | null = null;
      if (useGate) {
        gate = await buildGateNode(ctx, gateThreshold);
        head.connect(gate);
        head = gate;
      }
      const denoise = await buildModelNode(ctx, model);
      head.connect(denoise.node);

      const outAnalyser = ctx.createAnalyser();
      outAnalyser.fftSize = 1024;
      denoise.node.connect(outAnalyser);
      denoise.node.connect(ctx.destination); // loopback to speakers

      const tick = () => {
        setInDb(readDb(inAnalyser));
        setOutDb(readDb(outAnalyser));
        if (liveRef.current) liveRef.current.raf = requestAnimationFrame(tick);
      };
      liveRef.current = {
        ctx,
        stream,
        model: denoise,
        gate,
        raf: 0,
        inAnalyser,
        outAnalyser,
      };
      setLive(true);
      tick();
    } catch (e) {
      console.error('[denoise-tester] live monitor failed', e);
      stopLive();
    }
  };

  // ── Record & compare ─────────────────────────────────────────────────────
  const recRef = useRef<{
    ctx: AudioContext;
    stream: MediaStream;
    recorder: MediaRecorder;
    chunks: BlobPart[];
    raf: number;
    analyser: AnalyserNode;
    timer: number;
  } | null>(null);
  const clipRef = useRef<AudioBuffer | null>(null);
  const playRef = useRef<{ ctx: AudioContext; source: AudioBufferSourceNode } | null>(null);
  const [recording, setRecording] = useState(false);
  const [recDb, setRecDb] = useState(-100);
  const [hasClip, setHasClip] = useState(false);
  const [playing, setPlaying] = useState<string | null>(null);

  const teardownRecorder = () => {
    const r = recRef.current;
    recRef.current = null;
    if (r) {
      cancelAnimationFrame(r.raf);
      window.clearTimeout(r.timer);
      r.stream.getTracks().forEach((t) => t.stop());
      r.ctx.close().catch(() => undefined);
    }
  };

  const stopRecord = useCallback(() => {
    const r = recRef.current;
    if (r && r.recorder.state !== 'inactive') r.recorder.stop();
    setRecording(false);
  }, []);

  const startRecord = async () => {
    try {
      const stream = await navigator.mediaDevices.getUserMedia(RAW_CONSTRAINTS);
      const ctx = new AudioContext();
      const source = ctx.createMediaStreamSource(stream);
      const analyser = ctx.createAnalyser();
      analyser.fftSize = 1024;
      source.connect(analyser);

      const recorder = new MediaRecorder(stream);
      const chunks: BlobPart[] = [];
      recorder.ondataavailable = (ev) => {
        if (ev.data.size > 0) chunks.push(ev.data);
      };
      recorder.onstop = async () => {
        const blob = new Blob(chunks, { type: recorder.mimeType || 'audio/webm' });
        teardownRecorder();
        setRecDb(-100);
        try {
          const decodeCtx = new AudioContext({ sampleRate: 48000 });
          clipRef.current = await decodeCtx.decodeAudioData(await blob.arrayBuffer());
          decodeCtx.close().catch(() => undefined);
          setHasClip(true);
        } catch (e) {
          console.error('[denoise-tester] decode failed', e);
        }
      };

      const tick = () => {
        setRecDb(readDb(analyser));
        if (recRef.current) recRef.current.raf = requestAnimationFrame(tick);
      };
      const timer = window.setTimeout(stopRecord, MAX_RECORD_MS);
      recRef.current = { ctx, stream, recorder, chunks, raf: 0, analyser, timer };
      recorder.start();
      setRecording(true);
      tick();
    } catch (e) {
      console.error('[denoise-tester] record failed', e);
      teardownRecorder();
      setRecording(false);
    }
  };

  const stopPlayback = useCallback(() => {
    const p = playRef.current;
    playRef.current = null;
    if (p) {
      try {
        p.source.onended = null;
        p.source.stop();
      } catch {
        /* noop */
      }
      p.ctx.close().catch(() => undefined);
    }
    setPlaying(null);
  }, []);

  const play = async (label: string, playModel: DenoiseModelId | null) => {
    stopPlayback();
    const clip = clipRef.current;
    if (!clip) return;
    try {
      // bufferSource auto-resamples the 48 kHz clip to the context rate, so DTLN
      // gets the 16 kHz it needs while raw/RNNoise/Speex stay at 48 kHz.
      const ctx = new AudioContext({ sampleRate: sampleRateFor(playModel ?? 'rnnoise') });
      const source = ctx.createBufferSource();
      source.buffer = clip;
      if (playModel) {
        let head: AudioNode = source;
        if (useGate) {
          const gate = await buildGateNode(ctx, gateThreshold);
          head.connect(gate);
          head = gate;
        }
        const denoise = await buildModelNode(ctx, playModel);
        head.connect(denoise.node);
        denoise.node.connect(ctx.destination);
      } else {
        source.connect(ctx.destination);
      }
      source.onended = () => {
        if (playRef.current?.ctx === ctx) stopPlayback();
      };
      playRef.current = { ctx, source };
      source.start();
      setPlaying(label);
    } catch (e) {
      console.error('[denoise-tester] playback failed', e);
      stopPlayback();
    }
  };

  useEffect(
    () => () => {
      stopLive();
      stopPlayback();
      teardownRecorder();
    },
    [stopLive, stopPlayback],
  );

  const modelLabel = DENOISE_MODELS.find((m) => m.id === model)?.name ?? model;

  return (
    <Box direction="Column" gap="400">
      {/* Live monitor */}
      <Box direction="Column" gap="200">
        <Text size="T300">Live monitor</Text>
        <Text size="T200" priority="300">
          Hear yourself through {modelLabel}
          {useGate ? ' + gate' : ''} in real time. Use headphones to avoid feedback. Watch the In
          meter while you speak and while silent — set the gate threshold (orange line) just above
          your silent level.
        </Text>
        <Box direction="Row" gap="200" alignItems="Center">
          <Button
            size="300"
            variant={live ? 'Critical' : 'Secondary'}
            outlined
            onClick={live ? stopLive : startLive}
          >
            <Text size="T300">{live ? 'Stop' : 'Start monitor'}</Text>
          </Button>
        </Box>
        {live && (
          <Box direction="Column" gap="200">
            <DbMeter label="In (raw)" db={inDb} threshold={useGate ? gateThreshold : undefined} />
            <DbMeter label={`Out (${modelLabel})`} db={outDb} />
          </Box>
        )}
      </Box>

      {/* Record & compare */}
      <Box direction="Column" gap="200">
        <Text size="T300">Record &amp; compare</Text>
        <Text size="T200" priority="300">
          Record up to {MAX_RECORD_MS / 1000}s of yourself with your usual background noise, then
          play the same clip back raw vs through each model to A/B them. Captured fully raw (browser
          noise suppression off) so each model&apos;s effect is audible; uses the gate when enabled
          above.
        </Text>
        <Box direction="Row" gap="200" alignItems="Center">
          <Button
            size="300"
            variant={recording ? 'Critical' : 'Secondary'}
            outlined
            onClick={recording ? stopRecord : startRecord}
          >
            <Text size="T300">{recording ? 'Stop recording' : 'Record clip'}</Text>
          </Button>
          {recording && <Text size="T200">Recording…</Text>}
        </Box>
        {recording && <DbMeter label="Recording level" db={recDb} />}
        {hasClip && !recording && (
          <Box direction="Row" gap="200" wrap="Wrap">
            {(
              [
                { label: 'Raw', model: null },
                { label: 'RNNoise', model: 'rnnoise' },
                { label: 'Speex', model: 'speex' },
                { label: 'DTLN', model: 'dtln' },
                { label: 'DeepFilterNet', model: 'deepfilternet' },
              ] as const
            ).map((b) => (
              <Button
                key={b.label}
                size="300"
                variant={playing === b.label ? 'Primary' : 'Secondary'}
                outlined
                onClick={() => (playing === b.label ? stopPlayback() : play(b.label, b.model))}
              >
                <Text size="T300">
                  {playing === b.label ? `Stop ${b.label}` : `Play ${b.label}`}
                </Text>
              </Button>
            ))}
          </Box>
        )}
      </Box>
    </Box>
  );
}