#!/bin/bash
#
# Targeted Discovery Runner
# Executes specific discovery scripts on devices with partial data
#

set -euo pipefail

SCRIPT_DIR="$(dirname "$0")"

# Device configurations
declare -A PARTIAL_DEVICES
PARTIAL_DEVICES["fedora"]="localhost"
PARTIAL_DEVICES["lenovo420"]="100.98.144.95"
PARTIAL_DEVICES["lenovo"]="192.168.50.181"
PARTIAL_DEVICES["surface"]="100.67.40.97"

declare -A DEVICE_USERS
DEVICE_USERS["fedora"]="jonathan"
DEVICE_USERS["lenovo420"]="jon"
DEVICE_USERS["lenovo"]="jonathan"
DEVICE_USERS["surface"]="jon"

# Targeted scripts to run
SCRIPTS=(
    "targeted_security_discovery.sh"
    "targeted_data_discovery.sh"
    "targeted_performance_discovery.sh"
)

echo "=== Targeted Discovery Runner ==="
echo "Running missing discovery categories on partial devices"
echo "Devices: ${!PARTIAL_DEVICES[@]}"
echo "Scripts: ${SCRIPTS[@]}"
echo "======================================="

run_script_on_device() {
    local device=$1
    local host=${PARTIAL_DEVICES[$device]}
    local user=${DEVICE_USERS[$device]}
    local script=$2
    
    echo "[$device] Running $script"
    
    if [ "$host" = "localhost" ]; then
        # Local execution
        chmod +x "$SCRIPT_DIR/$script"
        sudo "$SCRIPT_DIR/$script"
    else
        # Remote execution
        scp "$SCRIPT_DIR/$script" "$user@$host:/tmp/"
        ssh "$user@$host" "chmod +x /tmp/$script && sudo /tmp/$script"
    fi
    
    echo "[$device] $script completed"
}

collect_results() {
    local device=$1
    local host=${PARTIAL_DEVICES[$device]}
    local user=${DEVICE_USERS[$device]}
    local results_dir="/home/jonathan/Coding/HomeAudit/targeted_discovery_results"
    
    mkdir -p "$results_dir"
    
    echo "[$device] Collecting results..."
    
    if [ "$host" = "localhost" ]; then
        # Local collection
        find /tmp -name "*_discovery_*_*" -type d -newer "$SCRIPT_DIR/$0" -exec cp -r {} "$results_dir/" \;
    else
        # Remote collection
        ssh "$user@$host" "find /tmp -name '*_discovery_*' -type d -newer /tmp/targeted_*_discovery.sh -exec tar -czf {}.tar.gz {} \;" 2>/dev/null || true
        scp "$user@$host:/tmp/*_discovery_*.tar.gz" "$results_dir/" 2>/dev/null || echo "[$device] No results to collect"
    fi
}

main() {
    local target_device=""
    local target_script=""
    
    # Parse command line arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --device)
                target_device="$2"
                shift 2
                ;;
            --script)
                target_script="$2"  
                shift 2
                ;;
            --help)
                echo "Usage: $0 [--device DEVICE] [--script SCRIPT]"
                echo "Devices: ${!PARTIAL_DEVICES[@]}"
                echo "Scripts: ${SCRIPTS[@]}"
                exit 0
                ;;
            *)
                echo "Unknown option: $1"
                exit 1
                ;;
        esac
    done
    
    # Run on specific device or all devices
    if [ -n "$target_device" ]; then
        if [[ ! " ${!PARTIAL_DEVICES[@]} " =~ " ${target_device} " ]]; then
            echo "Error: Unknown device '$target_device'"
            exit 1
        fi
        devices=("$target_device")
    else
        devices=("${!PARTIAL_DEVICES[@]}")
    fi
    
    # Run specific script or all scripts
    if [ -n "$target_script" ]; then
        if [[ ! " ${SCRIPTS[@]} " =~ " ${target_script} " ]]; then
            echo "Error: Unknown script '$target_script'"
            exit 1
        fi
        scripts=("$target_script")
    else
        scripts=("${SCRIPTS[@]}")
    fi
    
    # Execute targeted discovery
    for device in "${devices[@]}"; do
        echo "Starting targeted discovery on $device"
        
        for script in "${scripts[@]}"; do
            if ! run_script_on_device "$device" "$script"; then
                echo "Warning: $script failed on $device, continuing..."
            fi
            sleep 2  # Brief pause between scripts
        done
        
        collect_results "$device"
        echo "$device completed"
        echo "---"
    done
    
    echo "=== Targeted Discovery Complete ==="
    echo "Results available in: /home/jonathan/Coding/HomeAudit/targeted_discovery_results/"
    ls -la /home/jonathan/Coding/HomeAudit/targeted_discovery_results/ 2>/dev/null || echo "No results directory created yet"
}

main "$@"