#!/bin/bash
set -euo pipefail # Exit on error, undefined variable, or pipe failure

# Set STATE_FILE according to XDG Base Directory Specification
XDG_CACHE_HOME="${XDG_CACHE_HOME:-$HOME/.cache}"
readonly STATE_DIR="$XDG_CACHE_HOME/mute_toggle"
readonly STATE_FILE="$STATE_DIR/muted_sources"

# Store command paths or empty strings if not found
WPCTL_CMD=""
NOTIFY_SEND_CMD=""
CANBERRA_GTK_PLAY_CMD=""

# Function to play KDE notification sound
play_sound() {
    local event="$1"
    if [[ -n "$CANBERRA_GTK_PLAY_CMD" ]]; then
        "$CANBERRA_GTK_PLAY_CMD" -i "$event" --description="$event" --volume -10 2>/dev/null & disown
    fi
}

# Function to send desktop notification and play sound
# Arguments:
#   $1: Status message part (e.g., "Muted", "Unmuted")
#   $2: Icon name
#   $3: Sound event name
notify_user() {
    local status_msg="$1"
    local icon_name="$2"
    local sound_event="$3"

    play_sound "$sound_event"

    if [[ -n "$NOTIFY_SEND_CMD" ]]; then
        "$NOTIFY_SEND_CMD" -a "Microphone Toggle" -i "$icon_name" \
            "Microphone" "Microphone is $status_msg" -t 500
    fi
}

# Function to check for required commands
check_requirements() {
    if command -v wpctl &>/dev/null; then
        WPCTL_CMD="wpctl"
    else
        echo "Error: wpctl not found. Please install WirePlumber (package wireplumber)." >&2
        exit 1
    fi

    if command -v notify-send &>/dev/null; then
        NOTIFY_SEND_CMD="notify-send"
    else
        echo "Warning: notify-send not found. Desktop notifications will be skipped." >&2
    fi

    if command -v canberra-gtk-play &>/dev/null; then
        CANBERRA_GTK_PLAY_CMD="canberra-gtk-play"
    else
        echo "Warning: canberra-gtk-play not found. Notification sounds will be skipped." >&2
    fi
}

# Function to extract audio source IDs and their mute states
# Output: Lines of "<id> <0_or_1_for_muted_status>"
extract_sources() {
    # The awk script:
    # 1. /Audio/,/Video/: Operates only within the 'Audio' section up to 'Video'.
    # 2. If /Sources:/: Sets flag 's=1' indicating we are in the Sources sub-section.
    # 3. Else if /Filters:|Sinks:|Streams:/: Sets 's=0' if another sub-section starts.
    #    This handles cases where Sources is not the last sub-section.
    # 4. Else if (s && match(...)): If in Sources sub-section and line matches pattern for a source:
    #    - Extracts ID.
    #    - Prints ID and 1 if "MUTED]" is found, 0 otherwise.
    "$WPCTL_CMD" status 2>/dev/null | awk '
    /Audio/,/Video/ {
        if ($0 ~ /Sources:/) s=1
        else if ($0 ~ /Filters:|Sinks:|Streams:/) s=0
        else if (s && match($0, /^[[:space:]]*│[[:space:]]*[* ]?[[:space:]]*([0-9]+)\. (.*)/, a)) {
            id = a[1]
            print id, ($0 ~ /MUTED\]/ ? 1 : 0) # Check for "MUTED]" for more specificity
        }
    }'
}

main() {
    local FORCE_MODE=false
    if [[ "$#" -gt 0 && "$1" == "--force" ]]; then
        FORCE_MODE=true
        shift # Consume the --force argument
        echo "Running in --force mode: will not save/use saved state."
    fi

    check_requirements

    # Ensure state file directory exists (only if we might use it)
    if ! "$FORCE_MODE"; then
        mkdir -p "$STATE_DIR" || { echo "Error: Cannot create directory for $STATE_DIR" >&2; exit 1; }
    fi

    declare -A current_source_states  # Stores ID -> current_mute_status (0=unmuted, 1=muted)
    declare -A ids_muted_by_script    # Stores IDs of sources this script muted previously (from STATE_FILE)

    local -a ids_to_mute_now=()       # Array of UNMUTED source IDs to be muted
    local -a ids_to_unmute_based_on_script_state=() # Array of MUTED source IDs (that script muted) to be unmuted
    local -a ids_currently_muted=() # Array of ALL currently MUTED source IDs (used by --force)

    # Step 1: Load IDs of sources previously muted by this script (if not in force mode)
    if ! "$FORCE_MODE"; then
        if [[ -f "$STATE_FILE" ]]; then
            while IFS= read -r id_from_file || [[ -n "$id_from_file" ]]; do # Handle last line without newline
                if [[ "$id_from_file" =~ ^[0-9]+$ ]]; then # Basic validation
                    ids_muted_by_script["$id_from_file"]=1
                fi
            done < "$STATE_FILE"
        fi
    fi

    # Step 2: Get current state of all audio sources
    local no_sources_found=true
    while IFS=' ' read -r id muted_status; do
        no_sources_found=false
        current_source_states["$id"]=$muted_status
        if [[ "$muted_status" -eq 0 ]]; then # If currently unmuted
            ids_to_mute_now+=("$id")
        else # If currently muted
            ids_currently_muted+=("$id") # Store all muted IDs for --force unmute
            # If not in force mode, and it was muted by this script
            if ! "$FORCE_MODE" && [[ "${ids_muted_by_script[$id]+exists}" ]]; then
                ids_to_unmute_based_on_script_state+=("$id")
            fi
        fi
    done < <(extract_sources)

    if "$no_sources_found"; then
        echo "No active audio sources found."
        notify_user "Unavailable (No Sources)" "dialog-error" "dialog-error" # Or some other icon/sound
        exit 0
    fi

    # Step 3: Determine action based on force mode or state
    local target_ids=()
    local action_type="" # "mute" or "unmute"

    if "$FORCE_MODE"; then
        if [[ ${#ids_to_mute_now[@]} -gt 0 ]]; then # If any sources are unmuted, force mute
            target_ids=("${ids_to_mute_now[@]}")
            action_type="mute"
        elif [[ ${#ids_currently_muted[@]} -gt 0 ]]; then # If all are muted, force unmute
            target_ids=("${ids_currently_muted[@]}")
            action_type="unmute"
        fi
    else # Not force mode, use stateful logic
        if [[ ${#ids_to_mute_now[@]} -gt 0 ]]; then # If any sources are unmuted, mute them all
            target_ids=("${ids_to_mute_now[@]}")
            action_type="mute"
        elif [[ ${#ids_to_unmute_based_on_script_state[@]} -gt 0 ]]; then # If all are muted, and some by script, unmute those
            target_ids=("${ids_to_unmute_based_on_script_state[@]}")
            action_type="unmute"
        fi
    fi

    # Step 4: Perform the determined action
    if [[ "$action_type" == "mute" ]]; then
        echo "Muting sources: ${target_ids[*]}"
        if ! "$FORCE_MODE"; then # Clear state file if not in force mode
            >"$STATE_FILE" || { echo "Error: Cannot clear $STATE_FILE" >&2; exit 1; }
        fi
        for id in "${target_ids[@]}"; do
            if ! "$WPCTL_CMD" set-mute "$id" 1; then
                echo "Warning: Failed to mute source $id. Continuing..." >&2
            fi
            if ! "$FORCE_MODE"; then # Add to state file if not in force mode
                echo "$id" >> "$STATE_FILE"
            fi
        done
        notify_user "Muted" "microphone-sensitivity-muted" "dialog-warning"
    elif [[ "$action_type" == "unmute" ]]; then
        echo "Unmuting sources: ${target_ids[*]}"
        for id in "${target_ids[@]}"; do
            if ! "$WPCTL_CMD" set-mute "$id" 0; then
                 echo "Warning: Failed to unmute source $id. Continuing..." >&2
            fi
        done
        if ! "$FORCE_MODE"; then # Remove state file if not in force mode
            rm -f "$STATE_FILE" || { echo "Error: Cannot remove $STATE_FILE" >&2; exit 1; }
        fi
        notify_user "Unmuted" "microphone-sensitivity-high" "dialog-information"
    else
        # This case means:
        # 1. No sources are unmuted.
        # 2. AND (if not force mode) no sources currently muted were previously muted by this script.
        # This can happen if all sources are manually muted, or no relevant action was determined.
        echo "No microphone status change needed."
        if [[ -f "$STATE_FILE" ]]; then
             echo "(State file $STATE_FILE exists, but relevant sources are not currently active or already handled)."
        fi
        notify_user "Unchanged" "dialog-information" "dialog-information"
    fi
}

# Ensure script is not sourced
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi