#!/usr/bin/env bash
set -euo pipefail

log() { echo "[pulse-noise] $*"; }
die() { echo "[pulse-noise] ERROR: $*" >&2; exit 1; }

if [ "${EUID:-$(id -u)}" -ne 0 ]; then
  die "Run as root (or with sudo)."
fi

PULSE_DIR="/opt/pulse"
REPO_URL="https://github.com/plsechat/pulse-chat.git"

log "Preparing source in ${PULSE_DIR}..."
if [ -d "${PULSE_DIR}/.git" ]; then
  git -C "${PULSE_DIR}" fetch --all --prune
  git -C "${PULSE_DIR}" checkout .
  git -C "${PULSE_DIR}" pull --ff-only
else
  rm -rf "${PULSE_DIR}"
  git clone "${REPO_URL}" "${PULSE_DIR}"
fi

cd "${PULSE_DIR}"

# Ensure .env exists
if [ ! -f "${PULSE_DIR}/.env" ]; then
  log ".env missing ? creating from .env.supabase.example"
  cp -f .env.supabase.example .env

  if command -v bun >/dev/null 2>&1; then
    log "Generating Supabase keys..."
    KEY_OUTPUT="$(bun docker/generate-keys.ts)"
    JWT_SECRET="$(echo "${KEY_OUTPUT}" | grep -E '^JWT_SECRET=' | head -n1 | cut -d'=' -f2-)"
    SUPABASE_ANON_KEY="$(echo "${KEY_OUTPUT}" | grep -E '^SUPABASE_ANON_KEY=' | head -n1 | cut -d'=' -f2-)"
    SUPABASE_SERVICE_ROLE_KEY="$(echo "${KEY_OUTPUT}" | grep -E '^SUPABASE_SERVICE_ROLE_KEY=' | head -n1 | cut -d'=' -f2-)"

    sed -i \
      -e "s|^POSTGRES_PASSWORD=.*|POSTGRES_PASSWORD=$(openssl rand -hex 16)|" \
      -e "s|^JWT_SECRET=.*|JWT_SECRET=${JWT_SECRET}|" \
      -e "s|^SUPABASE_ANON_KEY=.*|SUPABASE_ANON_KEY=${SUPABASE_ANON_KEY}|" \
      -e "s|^SUPABASE_SERVICE_ROLE_KEY=.*|SUPABASE_SERVICE_ROLE_KEY=${SUPABASE_SERVICE_ROLE_KEY}|" \
      .env
  else
    log "bun not available; left .env defaults. Edit /opt/pulse/.env before running compose."
  fi
fi

log "Editing files directly..."
python3 <<'PY'
import json, re, pathlib

root = pathlib.Path("/opt/pulse")
def read(p): return p.read_text()
def write(p, s): p.write_text(s)

# package.json
pkg = root / "apps/client/package.json"
data = json.loads(read(pkg))
deps = data.setdefault("dependencies", {})
deps.setdefault("@timephy/rnnoise-wasm", "^1.0.0")
write(pkg, json.dumps(data, indent=2) + "\n")

# types.ts
types = root / "apps/client/src/types.ts"
t = read(types)
if "noiseSuppressionRnnoise" not in t:
    t = t.replace(
        "  noiseSuppression: boolean;\n",
        "  noiseSuppression: boolean;\n"
        "  noiseSuppressionEnhanced: boolean;\n"
        "  noiseSuppressionRnnoise: boolean;\n"
    )
    write(types, t)

# devices-provider defaults
devprov = root / "apps/client/src/components/devices-provider/index.tsx"
d = read(devprov)
if "noiseSuppressionRnnoise" not in d:
    d = d.replace(
        "  noiseSuppression: false,\n",
        "  noiseSuppression: false,\n"
        "  noiseSuppressionEnhanced: false,\n"
        "  noiseSuppressionRnnoise: false,\n"
    )
    write(devprov, d)

# settings UI
ui = root / "apps/client/src/components/server-screens/user-settings/devices/index.tsx"
u = read(ui)
if "RNNoise (WASM)" not in u:
    insert = """
            <Group label=\"Enhanced noise suppression\">
              <Switch
                checked={!!values.noiseSuppressionEnhanced}
                onCheckedChange={(checked) =>
                  onChange('noiseSuppressionEnhanced', checked)
                }
              />
            </Group>

            <Group label=\"RNNoise (WASM)\">
              <Switch
                checked={!!values.noiseSuppressionRnnoise}
                onCheckedChange={(checked) =>
                  onChange('noiseSuppressionRnnoise', checked)
                }
              />
            </Group>
"""
    marker = """            <Group label=\"Noise suppression\">
              <Switch
                checked={!!values.noiseSuppression}
                onCheckedChange={(checked) =>
                  onChange('noiseSuppression', checked)
                }
              />
            </Group>
"""
    if marker not in u:
        raise SystemExit("Could not find Noise suppression group.")
    u = u.replace(marker, marker + insert)
    write(ui, u)

# voice-provider
vp = root / "apps/client/src/components/voice-provider/index.tsx"
v = read(vp)

if "rnnoise-wasm" not in v:
    v = v.replace(
        "import { VolumeControlProvider } from './volume-control-context';\n",
        "import { VolumeControlProvider } from './volume-control-context';\n"
        "import { NoiseSuppressorWorklet_Name } from '@timephy/rnnoise-wasm';\n"
        "import NoiseSuppressorWorklet from '@timephy/rnnoise-wasm/NoiseSuppressorWorklet?worker&url';\n"
    )

if "buildMicConstraints" not in v:
    v = v.replace(
        "type AudioVideoRefs = {\n",
        """const buildMicConstraints = (devices: TDeviceSettings): MediaTrackConstraints => {
  const enhanced = !!devices.noiseSuppressionEnhanced;
  const rnnoise = !!devices.noiseSuppressionRnnoise;

  const constraints: MediaTrackConstraints = {
    deviceId: devices.microphoneId
      ? { exact: devices.microphoneId }
      : undefined,
    autoGainControl: (enhanced || rnnoise) ? true : devices.autoGainControl,
    echoCancellation: (enhanced || rnnoise) ? true : devices.echoCancellation,
    noiseSuppression: rnnoise ? false : (enhanced ? true : devices.noiseSuppression),
    sampleRate: 48000,
    channelCount: (enhanced || rnnoise) ? 1 : 2
  };

  if (enhanced) {
    const legacy = constraints as unknown as Record<string, unknown>;
    legacy.googHighpassFilter = true;
    legacy.googNoiseSuppression = true;
    legacy.googEchoCancellation = true;
    legacy.googAutoGainControl = true;
    legacy.googTypingNoiseDetection = true;
  }

  return constraints;
};

type AudioVideoRefs = {
"""
    )

if "applyRnnoise" not in v:
    v = v.replace(
        "} = useTransportStats();\n\n",
        """} = useTransportStats();

  type RNNoiseState = {
    ctx: AudioContext;
    source: MediaStreamAudioSourceNode;
    node: AudioWorkletNode;
    dest: MediaStreamAudioDestinationNode;
    input: MediaStream;
  };

  const rnnoiseRef = useRef<RNNoiseState | null>(null);
  const rawMicStreamRef = useRef<MediaStream | null>(null);

  const cleanupRnnoise = useCallback(() => {
    if (!rnnoiseRef.current) return;
    try {
      rnnoiseRef.current.node.disconnect();
      rnnoiseRef.current.source.disconnect();
      rnnoiseRef.current.dest.disconnect();
      rnnoiseRef.current.ctx.close();
    } catch {
      // ignore
    }
    rnnoiseRef.current = null;
  }, []);

  const stopRawMic = useCallback(() => {
    rawMicStreamRef.current?.getTracks().forEach((t) => t.stop());
    rawMicStreamRef.current = null;
  }, []);

  const applyRnnoise = useCallback(async (input: MediaStream): Promise<MediaStream> => {
    if (typeof AudioWorkletNode === 'undefined') {
      logVoice('RNNoise not supported: AudioWorklet unavailable');
      return input;
    }

    cleanupRnnoise();

    const ctx = new AudioContext({ sampleRate: 48000 });
    await ctx.audioWorklet.addModule(NoiseSuppressorWorklet);

    const source = ctx.createMediaStreamSource(input);
    const node = new AudioWorkletNode(ctx, NoiseSuppressorWorklet_Name);
    const dest = ctx.createMediaStreamDestination();

    source.connect(node);
    node.connect(dest);

    rnnoiseRef.current = { ctx, source, node, dest, input };

    return dest.stream;
  }, [cleanupRnnoise]);

  const acquireMicStream = useCallback(async (): Promise<{ stream: MediaStream; track: MediaStreamTrack; raw: MediaStream }> => {
    const raw = await navigator.mediaDevices.getUserMedia({
      audio: buildMicConstraints(devices),
      video: false
    });

    let stream = raw;
    if (devices.noiseSuppressionRnnoise) {
      try {
        stream = await applyRnnoise(raw);
      } catch (err) {
        logVoice('RNNoise failed, falling back to raw mic', { error: err });
        stream = raw;
      }
    }

    const track = stream.getAudioTracks()[0];
    return { stream, track, raw };
  }, [devices, applyRnnoise]);

"""
    )

v = re.sub(
    r"const stream = await navigator\.mediaDevices\.getUserMedia\(\{[\s\S]*?\}\);\n\n\s*logVoice\('Microphone stream obtained', \{ stream \}\);",
    "const { stream, track, raw } = await acquireMicStream();\n\n      logVoice('Microphone stream obtained', { stream });",
    v, count=1
)

v = v.replace("const audioTrack = stream.getAudioTracks()[0];\n\n", "")
v = v.replace("if (audioTrack) {\n        logVoice('Obtained audio track', { audioTrack });",
              "if (track) {\n        logVoice('Obtained audio track', { audioTrack: track });")
v = v.replace("track: audioTrack,\n", "track,\n")
v = v.replace("audioTrack.onended = () => {", "track.onended = () => {")

if "rawMicStreamRef.current = raw;" not in v:
    v = v.replace("setLocalAudioStream(stream);\n\n", "setLocalAudioStream(stream);\n\n      stopRawMic();\n      rawMicStreamRef.current = raw;\n\n")

if "cleanupRnnoise();" not in v:
    v = v.replace("localAudioProducer.current?.close();\n\n          setLocalAudioStream(undefined);",
                  "localAudioProducer.current?.close();\n\n          stopRawMic();\n          cleanupRnnoise();\n          setLocalAudioStream(undefined);")

v = v.replace(
    "devices.microphoneId,\n    devices.autoGainControl,\n    devices.echoCancellation,\n    devices.noiseSuppression",
    "acquireMicStream,\n    stopRawMic,\n    cleanupRnnoise"
)

v = re.sub(
    r"const newMicStream = await navigator\.mediaDevices\.getUserMedia\(\{[\s\S]*?\}\);\n\n\s*const newMicTrack",
    "const { stream: newMicStream, raw: newRaw } = await acquireMicStream();\n\n        const newMicTrack",
    v, count=1
)

if "rawMicStreamRef.current = newRaw;" not in v:
    v = v.replace("setLocalAudioStream(newMicStream);\n          logVoice('macOS: Mic restored to original settings');",
                  "setLocalAudioStream(newMicStream);\n          stopRawMic();\n          rawMicStreamRef.current = newRaw;\n          logVoice('macOS: Mic restored to original settings');")

v = v.replace(
    "devices.microphoneId, devices.autoGainControl, devices.echoCancellation, devices.noiseSuppression",
    "acquireMicStream,\n    stopRawMic"
)

v = re.sub(
    r"const newStream = await navigator\.mediaDevices\.getUserMedia\(\{[\s\S]*?\}\);\n\n\s*const newTrack = newStream\.getAudioTracks\(\)\[0\];",
    "const { stream: newStream, track: newTrack, raw: newRaw } =\n        await acquireMicStream();",
    v, count=1
)

if "rawMicStreamRef.current = newRaw;" not in v:
    v = v.replace("setLocalAudioStream(newStream);\n      logVoice('Mic settings reapplied successfully');",
                  "setLocalAudioStream(newStream);\n      stopRawMic();\n      rawMicStreamRef.current = newRaw;\n      logVoice('Mic settings reapplied successfully');")

v = v.replace(
    "devices.microphoneId,\n    devices.autoGainControl,\n    devices.echoCancellation,\n    devices.noiseSuppression",
    "acquireMicStream,\n    stopRawMic"
)

if "cleanupRnnoise();" not in v:
    v = v.replace("setConnectionStatus(ConnectionStatus.DISCONNECTED);\n  }, [",
                  "setConnectionStatus(ConnectionStatus.DISCONNECTED);\n    stopRawMic();\n    cleanupRnnoise();\n  }, [")
    v = v.replace("cleanupTransports", "cleanupTransports,\n    stopRawMic,\n    cleanupRnnoise")

v = v.replace(
    "prev.noiseSuppression !== devices.noiseSuppression ||\n      prev.autoGainControl !== devices.autoGainControl;",
    "prev.noiseSuppression !== devices.noiseSuppression ||\n      prev.autoGainControl !== devices.autoGainControl ||\n      prev.noiseSuppressionEnhanced !== devices.noiseSuppressionEnhanced ||\n      prev.noiseSuppressionRnnoise !== devices.noiseSuppressionRnnoise;"
)

write(vp, v)
PY

log "Building image pulse-chat:noise..."
docker build -t pulse-chat:noise .

log "Writing compose override..."
cat > docker-compose.noise.yml <<'YML'
services:
  pulse:
    image: pulse-chat:noise
YML

COMPOSE_FILES="-f docker-compose-supabase.yml"
if [ -f docker-compose.override.yml ]; then
  COMPOSE_FILES="$COMPOSE_FILES -f docker-compose.override.yml"
fi
COMPOSE_FILES="$COMPOSE_FILES -f docker-compose.noise.yml"

log "Restarting with noise build..."
docker compose $COMPOSE_FILES up -d --build

log "Done. UI now has 'Enhanced noise suppression' and 'RNNoise (WASM)' toggles."