#!/usr/bin/env bash
set -euo pipefail

# Script Name: network-discovery.sh
# Description: Discover devices on local network and highlight the newest device
# Version: 1.0.0
# Dependencies: arp-scan (or nmap), gum (optional but recommended)

# === Configuration ===
readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
readonly VERSION="1.0.0"
readonly LOGFILE="${LOGFILE:-/tmp/$(basename "$0" .sh)-$$.log}"

# Ensure log file is writable
touch "$LOGFILE" 2>/dev/null || LOGFILE="/dev/null"
chmod 644 "$LOGFILE" 2>/dev/null || true

# Colors for output
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[0;34m'
readonly CYAN='\033[0;36m'
readonly MAGENTA='\033[0;35m'
readonly BOLD='\033[1m'
readonly NC='\033[0m'

# === Logging Functions ===
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [INFO] $*" | tee -a "$LOGFILE"
}

log_error() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [ERROR] $*" | tee -a "$LOGFILE" >&2
}

log_warn() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [WARN] $*" | tee -a "$LOGFILE"
}

# === Cleanup Handler ===
TEMP_FILES=()

cleanup() {
    local exit_code=$?

    # Clean temp files
    for file in "${TEMP_FILES[@]}"; do
        [[ -f "$file" ]] && rm -f "$file"
    done

    # Clean log file on successful completion
    if [[ $exit_code -eq 0 ]] && [[ "$LOGFILE" != "/dev/null" ]]; then
        rm -f "$LOGFILE" 2>/dev/null || true
    fi

    exit $exit_code
}

trap cleanup EXIT INT TERM

# === Dependency Checking ===
HAS_GUM=false
HAS_ARP_SCAN=false
SCAN_METHOD=""

check_dependencies() {
    # Check for gum (optional) - check common locations
    if command -v gum &>/dev/null; then
        HAS_GUM=true
    elif [[ -x "$HOME/go/bin/gum" ]]; then
        HAS_GUM=true
        export PATH="$HOME/go/bin:$PATH"
    elif [[ -x "/home/e/go/bin/gum" ]]; then
        HAS_GUM=true
        export PATH="/home/e/go/bin:$PATH"
    fi

    # Check for scanning tools
    if command -v arp-scan &>/dev/null; then
        HAS_ARP_SCAN=true
        SCAN_METHOD="arp-scan"
    elif command -v nmap &>/dev/null; then
        SCAN_METHOD="nmap"
        log_warn "Using nmap (arp-scan recommended for better MAC detection)"
    else
        log_error "No network scanning tool found"
        echo "Please install one of:"
        echo "  sudo apt install arp-scan  (recommended)"
        echo "  sudo apt install nmap"
        return 1
    fi

    return 0
}

# === UI Functions ===
show_header() {
    clear

    if [[ "$HAS_GUM" == "true" ]]; then
        gum style \
            --border thick \
            --border-foreground 12 \
            --align center \
            --width 60 \
            --margin "1" \
            --padding "1 2" \
            "🔍 NETWORK DEVICE DISCOVERY" \
            "" \
            "v${VERSION}" \
            "Scanning local network..."
        echo
    else
        echo -e "${BLUE}╔════════════════════════════════════════════════════╗${NC}"
        echo -e "${BLUE}║${NC}     🔍 ${BOLD}NETWORK DEVICE DISCOVERY${NC}              ${BLUE}║${NC}"
        echo -e "${BLUE}║${NC}                                                    ${BLUE}║${NC}"
        echo -e "${BLUE}║${NC}               v${VERSION}                            ${BLUE}║${NC}"
        echo -e "${BLUE}║${NC}          Scanning local network...                 ${BLUE}║${NC}"
        echo -e "${BLUE}╚════════════════════════════════════════════════════╝${NC}"
        echo
    fi
}

# === Network Functions ===
get_local_network() {
    # Get the default gateway and derive network
    local gateway
    gateway=$(ip route | grep default | awk '{print $3}' | head -n1)

    if [[ -z "$gateway" ]]; then
        log_error "Could not determine default gateway"
        return 1
    fi

    # Extract network (assumes /24)
    local network
    network=$(echo "$gateway" | cut -d. -f1-3)
    echo "${network}.0/24"
}

scan_network_arp_scan() {
    local network="$1"
    local output_file="$2"

    if [[ "$HAS_GUM" == "true" ]]; then
        echo -e "${CYAN}🔍 Scanning network with arp-scan...${NC}"
        (
            sudo arp-scan --interface=eth0 --localnet 2>/dev/null || \
            sudo arp-scan --interface=wlan0 --localnet 2>/dev/null || \
            sudo arp-scan --localnet 2>/dev/null
        ) | tee "$output_file" &

        local scan_pid=$!
        # Spinner options: dot, pulse, points, minidot, line, jump, globe, moon, monkey, meter, hamburger
        gum spin --spinner pulse --title "Scanning local network..." -- bash -c "while kill -0 $scan_pid 2>/dev/null; do sleep 0.1; done"
        wait $scan_pid
    else
        echo -e "${CYAN}⏳ Scanning network with arp-scan...${NC}"
        sudo arp-scan --localnet 2>/dev/null | tee "$output_file"
    fi
}

scan_network_nmap() {
    local network="$1"
    local output_file="$2"

    if [[ "$HAS_GUM" == "true" ]]; then
        echo -e "${CYAN}🔍 Scanning network with nmap...${NC}"
        (
            sudo nmap -sn -PR "$network" 2>/dev/null
        ) | tee "$output_file" &

        local scan_pid=$!
        # Spinner options: dot, pulse, points, minidot, line, jump, globe, moon, monkey, meter, hamburger
        gum spin --spinner pulse --title "Scanning local network..." -- bash -c "while kill -0 $scan_pid 2>/dev/null; do sleep 0.1; done"
        wait $scan_pid
    else
        echo -e "${CYAN}⏳ Scanning network with nmap...${NC}"
        sudo nmap -sn -PR "$network" 2>/dev/null | tee "$output_file"
    fi
}

parse_arp_scan_results() {
    local scan_file="$1"
    local results_file="$2"

    # Parse arp-scan output: IP, MAC, Vendor
    # arp-scan format: 192.168.1.1  aa:bb:cc:dd:ee:ff  Vendor Name
    # Using pipe (|) as delimiter instead of comma to handle vendor names with commas
    grep -E "([0-9]{1,3}\.){3}[0-9]{1,3}" "$scan_file" | \
        grep -v "^Interface\|^Starting\|^Ending\|packets received" | \
        awk '{
            ip=$1
            mac=$2
            vendor=$3
            for(i=4;i<=NF;i++) vendor=vendor" "$i
            if(vendor=="") vendor="Unknown"
            print ip"|"mac"|"vendor
        }' > "$results_file"
}

parse_nmap_results() {
    local scan_file="$1"
    local results_file="$2"

    # After nmap scan, check entire ARP cache for all discovered devices
    log "Checking ARP cache for MAC addresses..."

    # Get all IPs from nmap output
    local found_ips=()
    while read -r line; do
        if [[ "$line" =~ "Nmap scan report for" ]]; then
            local ip=$(echo "$line" | grep -oE '([0-9]{1,3}\.){3}[0-9]{1,3}')
            [[ -n "$ip" ]] && found_ips+=("$ip")
        fi
    done < "$scan_file"

    # Now get MAC addresses from ARP cache
    for ip in "${found_ips[@]}"; do
        # Check arp cache
        local arp_line
        arp_line=$(arp -n | grep "^$ip " 2>/dev/null)

        if [[ -n "$arp_line" ]]; then
            # Parse: 10.98.0.1    ether   aa:bb:cc:dd:ee:ff   C   eth0
            local mac
            mac=$(echo "$arp_line" | awk '{print $3}')

            # Try to get vendor info (might need additional lookup)
            local vendor="Unknown"
            if [[ "$mac" =~ ^([0-9a-fA-F]{2}:){5}[0-9a-fA-F]{2}$ ]]; then
                # Valid MAC, try to identify device type
                case "${mac:0:8}" in
                    "00:50:56"|"00:0c:29"|"00:05:69") vendor="VMware" ;;
                    "08:00:27") vendor="VirtualBox" ;;
                    "52:54:00") vendor="QEMU/KVM" ;;
                    *) vendor="Device" ;;
                esac
            fi

            echo "${ip},${mac},${vendor}"
        else
            echo "${ip},Unknown,Unknown"
        fi
    done > "$results_file"
}

find_newest_device() {
    local results_file="$1"

    # Get current ARP cache with timestamps
    local newest_ip=""
    local newest_mac=""
    local newest_vendor=""
    local newest_age=999999

    # Read results and check ARP cache age (using pipe delimiter)
    while IFS='|' read -r ip mac vendor; do
        [[ -z "$ip" ]] && continue

        # Check if device is in ARP cache
        if arp -n "$ip" &>/dev/null; then
            # Most recently added device will be at the end of the list
            # We'll use the last device found as "newest"
            newest_ip="$ip"
            newest_mac="$mac"
            newest_vendor="$vendor"
        fi
    done < "$results_file"

    # If no ARP cache method works, just take the last device from scan
    if [[ -z "$newest_ip" ]]; then
        local last_line
        last_line=$(tail -n1 "$results_file")
        newest_ip=$(echo "$last_line" | cut -d'|' -f1)
        newest_mac=$(echo "$last_line" | cut -d'|' -f2)
        newest_vendor=$(echo "$last_line" | cut -d'|' -f3)
    fi

    echo "${newest_ip}|${newest_mac}|${newest_vendor}"
}

display_results() {
    local results_file="$1"
    local newest_device="$2"

    local newest_ip newest_mac newest_vendor
    IFS='|' read -r newest_ip newest_mac newest_vendor <<< "$newest_device"

    echo
    if [[ "$HAS_GUM" == "true" ]]; then
        gum style \
            --border double \
            --border-foreground 10 \
            --padding "1" \
            "📊 Discovered Devices"
        echo
    else
        echo -e "${GREEN}╔══════════════════════════════════════════════════╗${NC}"
        echo -e "${GREEN}║${NC}          📊 Discovered Devices                  ${GREEN}║${NC}"
        echo -e "${GREEN}╚══════════════════════════════════════════════════╝${NC}"
        echo
    fi

    # Header
    printf "${BOLD}%-16s %-20s %-30s${NC}\n" "IP ADDRESS" "MAC ADDRESS" "VENDOR"
    echo "────────────────────────────────────────────────────────────────────"

    # Display all devices - use awk to do ALL the formatting
    local device_count
    device_count=$(wc -l < "$results_file" 2>/dev/null || echo 0)

    # Use awk to format everything directly (avoids pipe/subshell issues)
    awk -F '|' -v newest_ip="$newest_ip" \
        -v MAGENTA="${MAGENTA}" -v CYAN="${CYAN}" -v YELLOW="${YELLOW}" \
        -v GREEN="${GREEN}" -v BOLD="${BOLD}" -v NC="${NC}" \
        '{
            ip=$1
            mac=$2
            vendor=$3

            if (ip == newest_ip) {
                # Newest device - HIGHLIGHT IT!
                printf "%s%s%-16s%s %s%s%-20s%s %s%-30s%s %s⭐ NEWEST%s\n", \
                    BOLD, MAGENTA, ip, NC, \
                    BOLD, CYAN, mac, NC, \
                    YELLOW, vendor, NC, \
                    GREEN, NC
            } else {
                printf "%-16s %-20s %-30s\n", ip, mac, vendor
            }
        }' "$results_file"

    echo
    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"

    # Summary box
    if [[ "$HAS_GUM" == "true" ]]; then
        gum style \
            --border rounded \
            --border-foreground 10 \
            --foreground 10 \
            --padding "1" \
            "✅ Scan Complete!" \
            "" \
            "Total devices found: ${device_count}" \
            "Newest device: ${newest_ip}" \
            "MAC Address: ${newest_mac}" \
            "Vendor: ${newest_vendor}"
    else
        echo
        echo -e "${GREEN}✅ Scan Complete!${NC}"
        echo -e "${BOLD}Total devices found:${NC} ${device_count}"
        echo
        echo -e "${BOLD}${MAGENTA}Newest Device:${NC}"
        echo -e "  ${BOLD}IP:${NC}     ${newest_ip}"
        echo -e "  ${BOLD}MAC:${NC}    ${newest_mac}"
        echo -e "  ${BOLD}Vendor:${NC} ${newest_vendor}"
    fi

    echo
}

# === Usage Function ===
usage() {
    cat << EOF
Usage: $(basename "$0") [OPTIONS]

Description:
    Scan local network for devices and highlight the newest device

Options:
    -h, --help      Show this help message
    -v, --verbose   Enable verbose output

Examples:
    sudo $(basename "$0")
    sudo $(basename "$0") --verbose

Requirements:
    - Must run with sudo (for network scanning)
    - arp-scan or nmap installed
    - gum (optional, for enhanced UI)

EOF
}

# === Main Logic ===
main() {
    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -h|--help)
                usage
                exit 0
                ;;
            -v|--verbose)
                set -x
                shift
                ;;
            *)
                log_error "Unknown option: $1"
                usage
                exit 1
                ;;
        esac
    done

    # Check if running as root
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root (for network scanning)"
        echo "Please run: sudo $0"
        exit 1
    fi

    # Check dependencies
    check_dependencies || exit 2

    # Show header
    show_header

    # Show arp-scan tip if using nmap
    if [[ "$SCAN_METHOD" == "nmap" ]]; then
        echo -e "${YELLOW}💡 Tip: Install arp-scan for better device detection${NC}"
        echo -e "${YELLOW}   Command: sudo apt install arp-scan${NC}"
        echo
    fi

    # Get local network
    log "Detecting local network..."
    local network
    network=$(get_local_network)
    log "Network: $network"
    echo -e "${BLUE}ℹ️  Network: ${BOLD}$network${NC}"
    echo

    # Create temp files
    local scan_file results_file
    scan_file=$(mktemp)
    results_file=$(mktemp)
    # Only add scan_file to cleanup - we need results_file until display is done
    TEMP_FILES+=("$scan_file")

    # Scan network
    log "Scanning network with $SCAN_METHOD"
    if [[ "$SCAN_METHOD" == "arp-scan" ]]; then
        scan_network_arp_scan "$network" "$scan_file"
        parse_arp_scan_results "$scan_file" "$results_file"
    else
        scan_network_nmap "$network" "$scan_file"
        parse_nmap_results "$scan_file" "$results_file"
    fi

    # Check if we found any devices
    if [[ ! -s "$results_file" ]]; then
        log_error "No devices found on network"
        exit 1
    fi

    # Find newest device
    log "Analyzing results..."
    local newest_device
    newest_device=$(find_newest_device "$results_file")

    # Display results
    display_results "$results_file" "$newest_device"

    # Clean up results file after display
    rm -f "$results_file"

    log "✅ Network discovery complete"
}

# Run main function
main "$@"