#!/usr/bin/env bash

# Training sample collection script for voice command classification
# Creates organized recordings for wake word and category phrases

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RECORDINGS_DIR="$SCRIPT_DIR/recordings"
SAMPLE_RATE=16000
DURATION=4

# Categories
declare -A CATEGORIES=(
    ["wake_word"]="ok, note this"
    ["pressure_below_15"]="eye pressure is less than 15"
    ["pressure_15_to_25"]="eye pressure is 15 to 25"
    ["pressure_25_to_35"]="eye pressure is 25 to 35"
    ["pressure_above_35"]="eye pressure is greater than 35"
    ["drop_l"]="L drop"
    ["drop_p"]="P drop"
    ["drop_d"]="D drop"
    ["both_eyes"]="both eyes"
)

# Get Corsair mic source
get_mic_source() {
    wpctl status | grep -i "corsair" | grep -i "mono" | grep -oE '[0-9]+\.' | head -1 | tr -d '.'
}

# Create directory structure
setup_dirs() {
    for category in "${!CATEGORIES[@]}"; do
        mkdir -p "$RECORDINGS_DIR/$category"
    done
    echo "Created recording directories in: $RECORDINGS_DIR"
}

# Count existing samples for a category
count_samples() {
    local category=$1
    find "$RECORDINGS_DIR/$category" -name "*.wav" 2>/dev/null | wc -l
}

# Record a single sample
record_sample() {
    local category=$1
    local voice_name=$2
    local source_id=$3

    local count=$(count_samples "$category")
    local filename="${voice_name}_$(printf "%03d" $((count + 1))).wav"
    local filepath="$RECORDINGS_DIR/$category/$filename"

    echo ""
    echo "Recording: $category"
    echo "Say: \"${CATEGORIES[$category]}\""
    echo ""
    echo "Press ENTER to start recording (${DURATION}s)..."
    read

    echo ">>> RECORDING NOW - SPEAK! <<<"
    pw-record --target "$source_id" --rate "$SAMPLE_RATE" --channels 1 --format s16 "$filepath" &
    local pid=$!
    sleep "$DURATION"
    kill "$pid" 2>/dev/null
    wait "$pid" 2>/dev/null

    # Check audio level
    local max_vol=$(ffmpeg -i "$filepath" -af "volumedetect" -f null /dev/null 2>&1 | grep "max_volume" | sed 's/.*max_volume: \([-0-9.]*\).*/\1/')

    echo "Saved: $filename (max volume: ${max_vol} dB)"

    if [ "${max_vol%.*}" -lt -50 ] 2>/dev/null; then
        echo "WARNING: Recording seems quiet. Keep? [Y/n]"
        read -r keep
        if [[ "$keep" =~ ^[Nn] ]]; then
            rm "$filepath"
            echo "Deleted. Try again."
            return 1
        fi
    fi
    return 0
}

# Interactive recording session
record_session() {
    local source_id=$(get_mic_source)

    if [ -z "$source_id" ]; then
        echo "ERROR: Corsair mic not found"
        exit 1
    fi

    echo "=== Voice Command Sample Collection ==="
    echo "Mic source: $source_id"
    echo ""

    echo "Enter your name/voice identifier (e.g., 'john', 'sarah'):"
    read -r voice_name
    voice_name=${voice_name:-default}

    while true; do
        echo ""
        echo "=== Current sample counts ==="
        for category in "${!CATEGORIES[@]}"; do
            local count=$(count_samples "$category")
            echo "  $category: $count samples"
        done

        echo ""
        echo "Select category to record:"
        echo "  1) wake_word           - \"ok, note this\""
        echo "  2) pressure_below_15   - \"eye pressure is less than 15\""
        echo "  3) pressure_15_to_25   - \"eye pressure is 15 to 25\""
        echo "  4) pressure_25_to_35   - \"eye pressure is 25 to 35\""
        echo "  5) pressure_above_35   - \"eye pressure is greater than 35\""
        echo "  6) drop_l              - \"L drop\""
        echo "  7) drop_p              - \"P drop\""
        echo "  8) drop_d              - \"D drop\""
        echo "  9) both_eyes           - \"both eyes\""
        echo "  a) Record ALL categories (one each)"
        echo "  q) Quit"
        echo ""
        read -r choice

        case $choice in
            1) record_sample "wake_word" "$voice_name" "$source_id" ;;
            2) record_sample "pressure_below_15" "$voice_name" "$source_id" ;;
            3) record_sample "pressure_15_to_25" "$voice_name" "$source_id" ;;
            4) record_sample "pressure_25_to_35" "$voice_name" "$source_id" ;;
            5) record_sample "pressure_above_35" "$voice_name" "$source_id" ;;
            6) record_sample "drop_l" "$voice_name" "$source_id" ;;
            7) record_sample "drop_p" "$voice_name" "$source_id" ;;
            8) record_sample "drop_d" "$voice_name" "$source_id" ;;
            9) record_sample "both_eyes" "$voice_name" "$source_id" ;;
            a|A)
                for category in wake_word pressure_below_15 pressure_15_to_25 pressure_25_to_35 pressure_above_35 drop_l drop_p drop_d both_eyes; do
                    record_sample "$category" "$voice_name" "$source_id"
                done
                ;;
            q|Q)
                echo "Done. Total samples collected:"
                for category in "${!CATEGORIES[@]}"; do
                    echo "  $category: $(count_samples "$category")"
                done
                exit 0
                ;;
            *) echo "Invalid choice" ;;
        esac
    done
}

# Main
setup_dirs
record_session