#!/bin/bash
#
# OMV-Optimized Discovery Script
# Optimized for OpenMediaVault systems - skips large data drives during migration
#

set -euo pipefail

# --- Configuration ---
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
HOSTNAME=$(hostname -f)
OUTPUT_BASE_DIR="/tmp/system_audit_${HOSTNAME}_${TIMESTAMP}"
DISCOVERY_DIR="${OUTPUT_BASE_DIR}/discovery"
mkdir -p "$DISCOVERY_DIR"
LOG_FILE="${OUTPUT_BASE_DIR}/discovery.log"

# OMV-specific exclusions for data drives that stay in place
OMV_DATA_PATHS=(
    "/srv/dev-disk-by-uuid-*"  # OMV data disks
    "/srv/mergerfs/*"          # MergerFS pools
    "/mnt/*"                   # External mounts
    "/media/*"                 # Media mounts
    "/export/*"                # NFS exports
)

# --- Logging ---
exec > >(tee -a "$LOG_FILE") 2>&1
echo "Starting OMV-Optimized State Discovery on ${HOSTNAME} at $(date)"
echo "Output will be saved in ${OUTPUT_BASE_DIR}"
echo "Excluding OMV data paths: ${OMV_DATA_PATHS[*]}"
echo "-----------------------------------------------------"

# --- Helper Functions ---
print_header() {
    echo ""
    echo "====================================================="
    echo ">= $1"
    echo "====================================================="
}

run_command() {
    local title="$1"
    local command="$2"
    local output_file="$3"
    
    print_header "$title"
    echo "Running command: $command"
    echo "Outputting to: $output_file"
    
    if eval "$command" > "$output_file"; then
        echo "Successfully captured $title."
    else
        echo "Warning: Command for '$title' failed or produced no output." > "$output_file"
    fi
}

# --- 1. Infrastructure Discovery ---
infra_discovery() {
    local out_dir="${DISCOVERY_DIR}/1_infrastructure"
    mkdir -p "$out_dir"
    
    run_command "CPU Information" "lscpu" "${out_dir}/cpu_info.txt"
    run_command "Memory Information" "free -h" "${out_dir}/memory_info.txt"
    run_command "PCI Devices (including GPU)" "lspci -v" "${out_dir}/pci_devices.txt"
    run_command "USB Devices" "lsusb -v" "${out_dir}/usb_devices.txt"
    run_command "Block Devices & Storage" "lsblk -o NAME,SIZE,TYPE,FSTYPE,MOUNTPOINT" "${out_dir}/storage_layout.txt"
    run_command "Filesystem Usage" "df -hT" "${out_dir}/disk_usage.txt"
    run_command "RAID Status" "cat /proc/mdstat || true" "${out_dir}/raid_status.txt"
    
    run_command "OS & Kernel Version" "cat /etc/os-release && uname -a" "${out_dir}/os_info.txt"
    run_command "Network Interfaces" "ip -br a" "${out_dir}/network_interfaces.txt"
    run_command "Routing Table" "ip r" "${out_dir}/routing_table.txt"
    run_command "DNS Configuration" "cat /etc/resolv.conf" "${out_dir}/dns_config.txt"
    run_command "Firewall Status (UFW)" "sudo ufw status verbose || true" "${out_dir}/firewall_ufw.txt"
    run_command "Firewall Status (iptables)" "sudo iptables -L -n -v || true" "${out_dir}/firewall_iptables.txt"
    
    # OMV-specific storage information
    run_command "OMV Storage Config" "omv-confdbadm read conf.system.storage.filesystem || true" "${out_dir}/omv_filesystems.txt"
    run_command "OMV Shares Config" "omv-confdbadm read conf.system.shares.sharedfolder || true" "${out_dir}/omv_shares.txt"
}

# --- 2. Services Inventory ---
services_inventory() {
    local out_dir="${DISCOVERY_DIR}/2_services"
    mkdir -p "$out_dir"
    
    # Docker
    if command -v docker &> /dev/null; then
        run_command "Docker Info" "docker info" "${out_dir}/docker_info.txt"
        run_command "Docker Running Containers" "docker ps -a" "${out_dir}/docker_ps.txt"
        run_command "Docker Images" "docker images" "${out_dir}/docker_images.txt"
        run_command "Docker Networks" "docker network ls" "${out_dir}/docker_networks.txt"
        run_command "Docker Volumes" "docker volume ls" "${out_dir}/docker_volumes.txt"
        
        print_header "Docker Container Details"
        for id in $(docker ps -q); do
            local name=$(docker inspect --format '{{.Name}}' "$id" | sed 's,^/,,')
            echo "Inspecting container: $name"
            docker inspect "$id" > "${out_dir}/container_${name}.json"
        done
        
        # OMV-Optimized Docker Compose Search - Skip data directories
        print_header "Finding Docker Compose files (OMV-optimized)"
        echo "Searching system directories only, excluding data drives..."
        
        # Build exclusion arguments for find command
        local exclude_args=""
        for path in "${OMV_DATA_PATHS[@]}"; do
            exclude_args="$exclude_args -path \"$path\" -prune -o"
        done
        
        # Search only essential system paths
        find /opt /home /etc /usr/local -maxdepth 5 -name "docker-compose.yml" -o -name "docker-compose.yaml" -o -name "compose.yml" > "${out_dir}/docker_compose_locations.txt" 2>/dev/null || true
        
        echo "Found $(wc -l < "${out_dir}/docker_compose_locations.txt") compose files"
        
        while IFS= read -r file; do
            if [ -f "$file" ]; then
                sudo cp "$file" "${out_dir}/compose_file_$(basename "$(dirname "$file")").yml" 2>/dev/null || true
            fi
        done < "${out_dir}/docker_compose_locations.txt"
        
        echo -e "\nContainer Management Tools:"
        docker ps --format "table {{.Names}}\t{{.Image}}\t{{.Ports}}" | grep -E "(portainer|watchtower|traefik|nginx-proxy|heimdall|dashboard)" > "${out_dir}/management_containers.txt" || echo "No common management tools detected" > "${out_dir}/management_containers.txt"
        
        echo -e "\nContainer Resource Usage:"
        docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.NetIO}}" 2>/dev/null | head -20 > "${out_dir}/container_resources.txt"
    else
        echo "Docker not found." > "${out_dir}/docker_status.txt"
    fi
    
    # Systemd Services
    run_command "Systemd Services (Enabled)" "systemctl list-unit-files --state=enabled" "${out_dir}/systemd_enabled_services.txt"
    run_command "Systemd Services (Running)" "systemctl list-units --type=service --state=running" "${out_dir}/systemd_running_services.txt"
    
    # OMV-specific services
    run_command "OMV Engine Status" "systemctl status openmediavault-engined || true" "${out_dir}/omv_engine_status.txt"
    run_command "OMV Web Interface Status" "systemctl status nginx || systemctl status apache2 || true" "${out_dir}/omv_web_status.txt"
}

# --- 3. Data & Storage Discovery (OMV-optimized) ---
data_discovery() {
    local out_dir="${DISCOVERY_DIR}/3_data_storage"
    mkdir -p "$out_dir"
    
    run_command "NFS Exports" "showmount -e localhost || true" "${out_dir}/nfs_exports.txt"
    run_command "Mounted File Systems" "mount" "${out_dir}/mounts.txt"
    run_command "Samba Shares" "smbstatus -S || true" "${out_dir}/samba_shares.txt"
    
    # OMV configuration exports
    print_header "OMV Configuration Backup"
    if command -v omv-confdbadm >/dev/null 2>&1; then
        omv-confdbadm read > "${out_dir}/omv_full_config.json" 2>/dev/null || echo "Could not export OMV config" > "${out_dir}/omv_config_error.txt"
        echo "OMV configuration backed up"
    fi
    
    print_header "Critical system directories only"
    # Skip data drives - only scan system paths
    find /etc /opt /usr/local -name "*.conf" -o -name "*.cfg" -o -name "*.ini" | head -100 > "${out_dir}/system_config_files.txt" 2>/dev/null || true
}

# --- 4. Security & Access Discovery ---
security_discovery() {
    local out_dir="${DISCOVERY_DIR}/4_security"
    mkdir -p "$out_dir"
    
    run_command "User Accounts" "cat /etc/passwd" "${out_dir}/users.txt"
    run_command "Sudoers Configuration" "sudo cat /etc/sudoers" "${out_dir}/sudoers.txt"
    run_command "SSH Daemon Configuration" "sudo cat /etc/ssh/sshd_config" "${out_dir}/sshd_config.txt"
    run_command "Last Logins" "last -a" "${out_dir}/last_logins.txt"
    run_command "Open Ports" "sudo ss -tuln" "${out_dir}/open_ports.txt"
    run_command "Cron Jobs (System)" "sudo cat /etc/crontab || true" "${out_dir}/crontab_system.txt"
    run_command "Cron Jobs (User)" "for user in \$(cut -f1 -d: /etc/passwd); do crontab -u \"\$user\" -l 2>/dev/null | sed \"s/^/[user] /\" ; done || true" "${out_dir}/crontab_users.txt"
}

# --- 5. Performance & Usage ---
performance_discovery() {
    local out_dir="${DISCOVERY_DIR}/5_performance"
    mkdir -p "$out_dir"
    
    run_command "Current Processes" "ps aux" "${out_dir}/processes.txt"
    run_command "Uptime & Load" "uptime" "${out_dir}/uptime.txt"
    run_command "Network Stats" "netstat -s" "${out_dir}/netstat.txt"
    run_command "IO Stats" "iostat -x 1 2" "${out_dir}/iostat.txt"
}

# --- Main Execution ---
main() {
    infra_discovery
    services_inventory
    data_discovery
    security_discovery
    performance_discovery
    
    print_header "Packaging Results"
    tar -czf "${OUTPUT_BASE_DIR}.tar.gz" -C "$(dirname "$OUTPUT_BASE_DIR")" "$(basename "$OUTPUT_BASE_DIR")"
    
    echo "-----------------------------------------------------"
    echo "OMV-Optimized Discovery complete."
    echo "Results packaged in ${OUTPUT_BASE_DIR}.tar.gz"
    echo "Data drives ($(echo "${OMV_DATA_PATHS[*]}" | tr ' ' ',')) were excluded from filesystem scan"
}

main