#!/bin/bash
#
# This file is part of PipeWire.
# SPDX-FileCopyrightText: Copyright © 2026 Chengyi Zhao <zhaochengyi@uniontech.com>
# SPDX-License-Identifier: MIT
#
# rnnoise-toggle — toggle the rnnoise noise-suppression source as the default.
#
# USAGE
#   rnnoise-toggle [-q|--quiet] on
#   rnnoise-toggle [-q|--quiet] off [<original-source>]
#   rnnoise-toggle [-q|--quiet] status
#   rnnoise-toggle -h|--help
#
# The caller may omit the <original-source> argument to "off"; the source
# saved by a prior "on" (in $XDG_RUNTIME_DIR) is then restored automatically.
# An explicit argument, if given, always takes precedence.
#
set -euo pipefail

readonly RNNOISE_SOURCE="effect_output.rnnoise"
readonly RNNOISE_INPUT="effect_input.rnnoise"
readonly STATE_FILE="${XDG_RUNTIME_DIR:-/run/user/${UID:-$(id -u)}}/rnnoise-toggle-original-source"

quiet=0

# --- output helpers: info -> stderr (silenced by -q), errors always stderr ---
info() { [[ "$quiet" -eq 0 ]] && printf '%s\n' "$*" >&2 || true; }
err()  { printf '%s\n' "$*" >&2; }

# --- audit logging (always logged to syslog, regardless of -q) --------------
# Logs every toggle action (on/off) with result (success/failure) and details.
audit_log() {
    local action="$1" result="$2" detail="${3:-}"
    logger -t rnnoise-toggle "user=${USER:-unknown} action=${action} result=${result} ${detail}" 2>/dev/null || true
}

# --- security helpers -------------------------------------------------------

# Validate that a source name contains only safe characters and is not the
# rnnoise source itself (prevents injection and self-referential restore).
validate_source_name() {
    local name="$1"
    if [[ -z "$name" ]]; then
        err "error: empty source name"
        return 1
    fi
    if [[ ! "$name" =~ ^[a-zA-Z0-9._-]+$ ]]; then
        err "error: invalid source name '$name' (only a-z A-Z 0-9 . _ - allowed)"
        return 1
    fi
    if [[ "$name" == "$RNNOISE_SOURCE" ]]; then
        err "error: cannot restore to the rnnoise source itself"
        return 1
    fi
    return 0
}

# Validate the runtime directory: must exist, be a directory, and be owned
# by the current user (prevents XDG_RUNTIME_DIR manipulation attacks).
validate_runtime_dir() {
    local dir
    dir="$(dirname "$STATE_FILE")"
    if [[ ! -d "$dir" ]]; then
        err "error: runtime directory '$dir' does not exist"
        exit 1
    fi
    local owner
    owner="$(stat -c '%u' "$dir" 2>/dev/null)" || {
        err "error: cannot stat runtime directory '$dir'"
        exit 1
    }
    if [[ "$owner" != "$(id -u)" ]]; then
        err "error: runtime directory '$dir' not owned by current user"
        exit 1
    fi
}

show_usage() {
    cat <<'EOF'
Usage: rnnoise-toggle [-q|--quiet] <command> [args]
       rnnoise-toggle -h|--help

Commands:
  on               Enable noise suppression (save original source, pin it as
                    effect_input's target, set default to effect_output.rnnoise)
   off [<source>]   Disable; restore <source>, or the saved original source
                    if <source> is omitted; unpin effect_input's target
   status           Print current state to stdout ("on" or "off")

Options:
  -q, --quiet      Suppress informational messages (errors still printed)
  -h, --help       Show this help

Exit codes:
  0  success
  1  usage error (bad/missing arguments)
  2  dependency missing (pactl or pw-metadata not in PATH)
  3  command failed (rnnoise source not available, or set-default failed)
EOF
}

# --- dependency check -------------------------------------------------------
missing_deps=()
if ! command -v pactl >/dev/null 2>&1; then
    missing_deps+=("pactl")
fi
if ! command -v pw-metadata >/dev/null 2>&1; then
    missing_deps+=("pw-metadata")
fi
if [[ ${#missing_deps[@]} -gt 0 ]]; then
    err "error: missing dependencies: ${missing_deps[*]}"
    err "       (is pipewire-utils / pulseaudio-utils installed?)"
    exit 2
fi

# Validate runtime directory before any file operations
validate_runtime_dir

# --- helpers ----------------------------------------------------------------
check_rnnoise_source() {
    if ! pactl list sources short 2>/dev/null | grep -qw "$RNNOISE_SOURCE"; then
        err "error: rnnoise source '$RNNOISE_SOURCE' not found."
        err "       Is the filter-chain configuration loaded?"
        exit 3
    fi
}

# Get the PipeWire object.serial for a node by its node.name.
# Prints the serial number on stdout, or nothing if not found.
get_node_serial() {
    local name="$1"
    pw-cli ls Node 2>/dev/null \
        | awk -v n="$name" '
            $0 ~ /^\s*id [0-9]+/ { id=$2; sub(/,/, "", id) }
            $0 ~ "node.name = \"" n "\"" { print id; exit }
        '
}

# Pin effect_input.rnnoise to a specific source via pw-metadata target.object.
# This prevents WirePlumber from re-routing to a different source when the
# default source changes (e.g. falling back to a headset instead of the mic).
pin_input_source() {
    local target_source="$1"
    local serial
    serial="$(get_node_serial "$RNNOISE_INPUT")" || true
    if [[ -z "$serial" ]]; then
        err "warning: could not find '$RNNOISE_INPUT' in PipeWire; skipping target pin"
        return 0
    fi
    if ! pw-metadata -n default "$serial" target.object "$target_source" >/dev/null 2>&1; then
        err "warning: failed to pin target.object on '$RNNOISE_INPUT'"
        return 0
    fi
    info "pinned effect_input to: $target_source"
}

# Unpin effect_input.rnnoise target.object (let WirePlumber auto-route again).
unpin_input_source() {
    local serial
    serial="$(get_node_serial "$RNNOISE_INPUT")" || true
    if [[ -z "$serial" ]]; then
        return 0
    fi
    pw-metadata -n default -d "$serial" target.object >/dev/null 2>&1 || true
}

# Save the current default source to the state file.
# Returns 0 on success, 1 on failure (does not exit).
save_original_source() {
    local current
    current="$(pactl get-default-source 2>/dev/null)" || true
    # If no default source exists, there is nothing to save (this is rare).
    [[ -z "$current" ]] && return 0
    validate_source_name "$current" || return 1

    # Security: refuse to write through a symlink (prevents symlink attack)
    if [[ -L "$STATE_FILE" ]]; then
        err "error: state file '$STATE_FILE' is a symlink, refusing to write"
        audit_log "on" "failure" "state_file_symlink"
        return 1
    fi

    # Security: atomic write via temp file + rename (prevents partial write)
    local tmp="${STATE_FILE}.$$"
    printf '%s\n' "$current" > "$tmp" || {
        err "error: failed to write state file"
        rm -f "$tmp"
        return 1
    }
    chmod 600 "$tmp"
    mv -f "$tmp" "$STATE_FILE"
    return 0
}

# Read the saved original source from the state file.
# Returns:
#   0 and prints the source name if file exists and content is valid.
#   0 (empty) if file does not exist (no saved source).
#   1 and prints an error if file exists but content is invalid or symlink.
read_original_source() {
    [[ -f "$STATE_FILE" ]] || return 0
    # Security: refuse to read through a symlink
    if [[ -L "$STATE_FILE" ]]; then
        err "error: state file '$STATE_FILE' is a symlink, refusing to read"
        return 1
    fi
    local content
    content="$(cat "$STATE_FILE" 2>/dev/null)" || return 0
    # Security: validate content before use
    if ! validate_source_name "$content" 2>/dev/null; then
        err "error: state file content is invalid, ignoring saved source"
        return 1
    fi
    printf '%s\n' "$content"
}

# List available sources (excluding the rnnoise source) to stderr.
list_available_sources() {
    local name
    pactl list sources short 2>/dev/null | awk '{print $2}' | while read -r name; do
        if [[ "$name" != "$RNNOISE_SOURCE" ]]; then
            err "  $name"
        fi
    done || true
}

# --- argument parsing -------------------------------------------------------
cmd=""
src=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        -q|--quiet)
            quiet=1; shift ;;
        -h|--help)
            show_usage; exit 0 ;;
        --)
            shift; break ;;
        on|enable|start)
            [[ -n "$cmd" ]] && { err "error: multiple commands given"; show_usage >&2; exit 1; }
            cmd="on"; shift ;;
        off|disable|stop)
            [[ -n "$cmd" ]] && { err "error: multiple commands given"; show_usage >&2; exit 1; }
            cmd="off"; shift
            if [[ $# -gt 0 && "$1" != -* ]]; then
                src="$1"; shift
            fi ;;
        status)
            [[ -n "$cmd" ]] && { err "error: multiple commands given"; show_usage >&2; exit 1; }
            cmd="status"; shift ;;
        *)
            err "error: unknown argument '$1'"
            show_usage >&2; exit 1 ;;
    esac
done

if [[ -z "$cmd" ]]; then
    err "error: no command given"
    show_usage >&2; exit 1
fi

# --- commands ---------------------------------------------------------------
do_on() {
    check_rnnoise_source

    # Save current default source BEFORE switching.
    local original_source
    original_source="$(pactl get-default-source 2>/dev/null)" || true

    if ! save_original_source; then
        audit_log "on" "failure" "save_source_failed"
        exit 1
    fi

    # Pin effect_input.rnnoise to the original source so WirePlumber won't
    # re-route it to a different mic when we change the default.
    if [[ -n "$original_source" ]]; then
        validate_source_name "$original_source" 2>/dev/null && \
            pin_input_source "$original_source"
    fi

    if ! pactl set-default-source "$RNNOISE_SOURCE"; then
        err "error: failed to set default source to '$RNNOISE_SOURCE'"
        audit_log "on" "failure" "set_default_failed"
        exit 3
    fi
    info "rnnoise ON (default source: $RNNOISE_SOURCE)"
    audit_log "on" "success" "source=$RNNOISE_SOURCE input_target=${original_source:-auto}"
}

do_off() {
    local target="$src"
    if [[ -z "$target" ]]; then
        # Try to read saved source; if read fails (invalid content), exit.
        if ! target="$(read_original_source)"; then
            err "Failed to read saved source. Available sources:"
            list_available_sources
            exit 1
        fi
        if [[ -z "$target" ]]; then
            err "error: no original source saved and no source specified."
            err ""
            err "Available sources (use with: rnnoise-toggle off <name>):"
            list_available_sources
            err ""
            err "Usage: rnnoise-toggle off <source-name>"
            exit 1
        fi
    fi
    # Security: validate user-supplied or file-read source name
    if ! validate_source_name "$target"; then
        audit_log "off" "failure" "invalid_source_name"
        exit 1
    fi

    # Unpin effect_input.rnnoise so WirePlumber can auto-route again.
    unpin_input_source

    if ! pactl set-default-source "$target"; then
        err "error: failed to set default source to '$target'"
        audit_log "off" "failure" "set_default_failed target=$target"
        exit 3
    fi
    info "rnnoise OFF (restored source: $target)"
    audit_log "off" "success" "source=$target"
}

do_status() {
    local current
    current="$(pactl get-default-source 2>/dev/null)" || current=""
    if [[ -z "$current" ]]; then
        # pactl failed — do not print "off" (which would be misleading)
        err "error: cannot query default source (is pipewire-pulse running?)"
        exit 3
    fi
    if [[ "$current" = "$RNNOISE_SOURCE" ]]; then
        printf 'on\n'
    else
        printf 'off\n'
    fi
}

case "$cmd" in
    on)      do_on ;;
    off)     do_off ;;
    status)  do_status ;;
esac

exit 0
