#!/bin/bash
# Toggle dictation: 1st run starts recording, 2nd run stops, transcribes locally
# with whisper.cpp, copies the text to the clipboard (Klipper history), appends it
# to ~/.local/share/dictation-history.txt, and types it into the focused window.
#
#   dictate            toggle
#   dictate --status   print idle / recording / transcribing
#
# Settings in ~/.config/frame-tweaks/dictate.conf: DICTATE_MODEL, DICTATE_LANG (default en), DICTATE_PROMPT,
# DICTATE_NOTYPE=1 (copy only), DICTATE_MAX_SECONDS (auto-stop, default 300).
export PATH="$HOME/.local/bin:/home/linuxbrew/.linuxbrew/bin:$PATH"
# Translations (lib/i18n.sh; the copy install_i18n puts here), else English
# shellcheck source=/dev/null
source "$HOME/.local/share/frame-tweaks/i18n/i18n.sh" 2>/dev/null ||
    t() { local f=$1; shift; if [ $# -eq 0 ]; then printf '%s' "$f"; else printf -- "$f" "$@"; fi; }
RUN=/run/user/$(id -u)
export YDOTOOL_SOCKET=$RUN/.ydotool_socket
# Shortcut/launcher runs get XDG_RUNTIME_DIR=/run/user/<uid>, but the desktop's
# Wayland socket lives under nested_plasma/ — point wl-copy at it explicitly.
if [ ! -S "${XDG_RUNTIME_DIR:-$RUN}/${WAYLAND_DISPLAY:-wayland-0}" ]; then
    for s in "$RUN"/nested_plasma/wayland-0 "$RUN"/*/wayland-0; do
        [ -S "$s" ] && { export WAYLAND_DISPLAY="$s"; break; }
    done
fi
CONF="$HOME/.config/frame-tweaks/dictate.conf"
# (before 0.6 it was ~/.config/dictate.conf: moved here once)
if [ ! -f "$CONF" ] && [ -f "$HOME/.config/dictate.conf" ]; then
    mkdir -p "${CONF%/*}" && mv "$HOME/.config/dictate.conf" "$CONF"
fi
# shellcheck source=/dev/null
[ -f "$CONF" ] && . "$CONF"
# Same model file as whisper-frame (groffta/whisper-frame), so the two never keep two copies.
MODEL="${DICTATE_MODEL:-$HOME/.local/share/whisper/ggml-small-q8_0.bin}"
LANGUAGE="${DICTATE_LANG:-en}"     # the shared model is multilingual: say which language
# whisper-frame bundles a GPU build of whisper-cli; use it when present, else Homebrew's.
WHISPER="$HOME/.local/lib/whisper-frame/whisper-cli"
[ -x "$WHISPER" ] || WHISPER=$(command -v whisper-cli)
# Vocabulary hint so Whisper spells these words correctly.
PROMPT="${DICTATE_PROMPT:-SteamOS, Steam Frame, SteamVR, KDE Plasma.}"
MAX_SECONDS="${DICTATE_MAX_SECONDS:-300}"
# Fixed path: launcher, shortcut and terminal see different XDG_RUNTIME_DIRs.
STATE=$RUN/dictate
HISTORY="$HOME/.local/share/dictation-history.txt"
LOG="$STATE/dictate.log"
SOUNDS=/usr/share/sounds/ocean/stereo
mkdir -p "$STATE"
PID="$STATE/rec.pid"        # pid of the running pw-record
CUR="$STATE/rec.current"    # path of the file being recorded
BUSY="$STATE/busy.pid"      # pid of a running transcription
LAST="$STATE/last-toggle"   # time of the last toggle, to ignore double taps
TIMER="$STATE/autostop.pid" # pid of the recording's auto-stop timer, ended with the recording

notify() { notify-send -a Dictation -t 1500 "$@" 2>/dev/null; }
cue() { pw-play "$SOUNDS/$1.oga" 9>&- 2>/dev/null & }   # 9>&-: never hold the toggle lock
log() { echo "$(date +%T) $*" >> "$LOG"; }
alive() { [ -f "$1" ] && kill -0 "$(cat "$1")" 2>/dev/null; }

# type_text TEXT: into the focused window. ydotool's own typing assumes a US keyboard, so on
# another layout (KDE's, in desktop mode) the keys for each character come from that layout
# (xkb.py, beside i18n.sh); text with a character the layout has no key for isn't typed at all
# (it's in the clipboard). Fails (1) when nothing was typed.
# A modifier held down on ydotoold's keyboard (Arrow Pad's Ctrl, latched or locked) would turn the
# text into shortcuts (Ctrl+W closes the tab, Ctrl+Q quits), so those are let go first (xkb.py let-go:
# only the ones down; Alt and Super around F24, so that's no tap opening a menu).
type_text() {
    local out rc=0 keys mods xkb="${FRAME_TWEAKS_PYLIB:-$HOME/.local/share/frame-tweaks/i18n}/xkb.py"
    out=$(python3 "$xkb" ydotool "$1" 2>>"$LOG") || rc=$?
    if [ $rc = 0 ] || [ $rc = 3 ]; then
        read -ra mods <<<"$(python3 "$xkb" let-go 2>>"$LOG")"
        if [ ${#mods[@]} -gt 0 ]; then
            log "letting go of modifiers: ${mods[*]}"
            ydotool key --key-delay 11 "${mods[@]}" 2>>"$LOG"
        fi
    fi
    case $rc in
        3) ydotool type --key-delay 2 -- "$1" 2>>"$LOG" || log "typing failed (is ydotoold running?)" ;;
        0) read -ra keys <<<"$out"         # (CODE:1 / CODE:0, a key down / up, with Shift or AltGr around)
           # 11 ms between key events: about ydotool type's pace (20 ms held, 2 between keys)
           ydotool key --key-delay 11 "${keys[@]}" 2>>"$LOG" || log "typing failed (is ydotoold running?)" ;;
        2) log "not typed: this keyboard layout has no key for $out"
           notify "$(t 'Not typed: copied to the clipboard')" \
               "$(t 'Your keyboard layout has no key for %s. Paste the text instead (Ctrl+V).' "$out")"
           return 1 ;;
        *) log "not typed: this keyboard layout couldn't be read"
           notify "$(t 'Not typed: copied to the clipboard')" "$(t 'Paste the text instead (Ctrl+V).')"
           return 1 ;;
    esac
}

if [ "${1:-}" = --status ]; then
    if alive "$PID"; then echo recording; elif alive "$BUSY"; then echo transcribing; else echo idle; fi
    exit 0
fi

transcribe() {   # transcribe <wav>: runs after the lock is released
    local wav=$1 text
    echo $$ > "$BUSY"
    cue device-removed
    notify "$(t 'Transcribing…')"
    log "stop: $(ffprobe -v error -show_entries format=duration -of csv=p=0 "$wav")s of audio"
    # Whisper invents text for silence ("Thanks for watching", "Subs by …"), so check
    # loudness first: a recording whose peak stays under -40 dB has no speech.
    local peak
    peak=$(ffmpeg -nostats -i "$wav" -af volumedetect -f null - 2>&1 | sed -n 's/.*max_volume: \(-*[0-9.]*\) dB/\1/p')
    if [ -n "$peak" ] && awk -v p="$peak" 'BEGIN{exit !(p < -40)}'; then
        log "no speech (peak ${peak} dB)"; mv -f "$wav" "$STATE/last.wav"; rm -f "$BUSY"
        notify "$(t 'No speech detected')"; return
    fi
    # With "Memory caps" installed, Whisper runs inside its memory limit.
    local cap=(); [ -x "$HOME/.local/bin/capped" ] && cap=("$HOME/.local/bin/capped" whisper --)
    text=$(DISABLE_GAMESCOPE_WSI=1 "${cap[@]}" "$WHISPER" -m "$MODEL" -f "$wav" -t 5 -nt -np -l "$LANGUAGE" \
            --prompt "$PROMPT" 2>>"$LOG" \
        | tr '\n' ' ' | sed -E 's/\[[A-Z_ ]+\]//g; s/\([a-z ]+\)//g; s/^ +//; s/ +$//; s/  +/ /g')
    # Drop Whisper's well-known phantom phrases when they are the whole result.
    if grep -qiE '^(subs? (by|from) .*|subtitles by .*|thanks? (you )?for watching[.!]?|thank you[.!]?|you|bye[.!]?|\.+)$' <<<"$text"; then
        log "dropped likely hallucination: $text"; text=""
    fi
    mv -f "$wav" "$STATE/last.wav"
    rm -f "$BUSY"
    if [ -z "$text" ]; then log "no speech"; notify "$(t 'No speech detected')"; return; fi
    # Trailing space after a finished sentence, so the next dictation flows on.
    [[ "$text" =~ [.?!]$ ]] && text="$text "
    printf '%s\t%s\n' "$(date '+%F %T')" "$text" >> "$HISTORY"
    log "text: $text"
    printf '%s' "$text" | setsid wl-copy 2>>"$LOG" || log "wl-copy failed (WAYLAND_DISPLAY=$WAYLAND_DISPLAY)"
    if [ -z "${DICTATE_NOTYPE:-}" ]; then
        type_text "$text" || return 0          # (it said so)
    fi
    notify "$(t 'Copied to clipboard')" "$text"
}

# Decide start/stop under a lock, so two quick presses can't both start (or both stop).
exec 9>"$STATE/lock"
flock 9
now=$(date +%s%N)
if [ -f "$LAST" ] && [ $(( (now - $(cat "$LAST")) / 1000000 )) -lt 700 ]; then
    log "ignored double tap"; exit 0
fi
echo "$now" > "$LAST"

if [ -f "$PID" ]; then
    # Stop, and transcribe. (A recording whose pw-record ended by itself, the microphone gone say,
    # counts as one still running: what it caught is transcribed, not left behind for a new one.)
    wav=$(cat "$CUR")
    if alive "$PID"; then kill -INT "$(cat "$PID")"; else log "the recording had ended by itself"; fi
    rm -f "$PID"
    # Its auto-stop timer goes too (it slept on for MAX_SECONDS; only if it's still a dictate).
    t=$(cat "$TIMER" 2>/dev/null) && grep -qs dictate "/proc/$t/cmdline" && kill "$t" 2>/dev/null
    rm -f "$TIMER"
    echo $$ > "$BUSY"                 # mark busy before unlocking
    flock -u 9; exec 9>&-
    sleep 0.3
    transcribe "$wav"
elif alive "$BUSY"; then
    log "ignored: still transcribing"
    notify "$(t 'Still transcribing…')" "$(t 'Wait for the text before dictating again.')"
else
    wav="$STATE/rec-$(date +%s%N).wav"   # own file per recording: never clobbers one being transcribed
    pw-record --rate 16000 --channels 1 --format s16 "$wav" 9>&- 2>/dev/null &   # 9>&-: don't inherit the lock
    rec=$!
    echo "$rec" > "$PID"; echo "$wav" > "$CUR"
    log "start pid $rec"
    # Safety net: stop a forgotten recording after MAX_SECONDS (if it's still this one: the toggle
    # then stops and transcribes it, or, if pw-record ended by itself, transcribes what it caught).
    # (Its sleep in the background, so that the stop above, killing it, ends the sleep too.)
    ( exec 9>&-; trap 'kill "$s" 2>/dev/null; exit 0' TERM
      sleep "$MAX_SECONDS" & s=$!; wait "$s"
      if [ "$(cat "$PID" 2>/dev/null)" = "$rec" ]; then
          log "auto-stop after ${MAX_SECONDS}s"; rm -f "$LAST" "$TIMER"; exec "$0"     # ($TIMER: that's us)
      fi
    ) >/dev/null 2>&1 &
    echo $! > "$TIMER"
    sleep 0.15; cue device-added
    notify "🎙 $(t 'Listening… (run again to stop)')"
fi
