# Windows Network Scanner for WOL Manager
# This script scans the local network for devices and extracts their MAC addresses
# Results are saved to a JSON file that can be read by the Docker container

param(
    [string]$Network = "",
    [string]$OutputPath = ".\data\network-scan-results.json",
    [int]$TimeoutMs = 1000,
    [int]$BatchSize = 10,
    [switch]$Verbose
)

function Write-Log {
    param([string]$Message)
    if ($Verbose) {
        Write-Host "[$(Get-Date -Format 'HH:mm:ss')] $Message" -ForegroundColor Green
    }
}

function Get-LocalNetworks {
    Write-Log "Detecting local networks..."
    $networks = @()
    
    # Get all network adapters that are up and have an IP address
    $adapters = Get-NetIPAddress -AddressFamily IPv4 | Where-Object { 
        $_.IPAddress -ne "127.0.0.1" -and ($_.PrefixOrigin -eq "Dhcp" -or $_.PrefixOrigin -eq "Manual")
    }
    
    foreach ($adapter in $adapters) {
        $networkAddr = $adapter.IPAddress
        $prefixLength = $adapter.PrefixLength
        
        # Network calculation for different prefix lengths
        $parts = $networkAddr.Split('.')
        if ($parts.Length -eq 4) {
            if ($prefixLength -eq 24) {
                $networkAddress = "$($parts[0]).$($parts[1]).$($parts[2]).0"
                $networks += "$networkAddress/24"
                Write-Log "Found network: $networkAddress/24 from adapter $($adapter.IPAddress)"
            }
            elseif ($prefixLength -eq 20) {
                # For /20, calculate the network base (e.g., 10.0.20.144/20 -> 10.0.16.0/20)
                $thirdOctet = [int]$parts[2]
                $networkThird = $thirdOctet - ($thirdOctet % 16)  # Round down to nearest /20 boundary
                $networkAddress = "$($parts[0]).$($parts[1]).$networkThird.0"
                $networks += "$networkAddress/20"
                Write-Log "Found network: $networkAddress/20 from adapter $($adapter.IPAddress)"
            }
            elseif ($prefixLength -eq 16) {
                $networkAddress = "$($parts[0]).$($parts[1]).0.0"
                $networks += "$networkAddress/16"
                Write-Log "Found network: $networkAddress/16 from adapter $($adapter.IPAddress)"
            }
            else {
                Write-Log "Unsupported prefix length /$prefixLength for $networkAddr, trying /24"
                # Fallback to /24 for unknown prefix lengths
                $networkAddress = "$($parts[0]).$($parts[1]).$($parts[2]).0"
                $networks += "$networkAddress/24"
                Write-Log "Fallback network: $networkAddress/24 from adapter $($adapter.IPAddress)"
            }
        }
    }
    
    # Add common default networks if none found
    if ($networks.Count -eq 0) {
        $networks += "192.168.1.0/24"
        Write-Log "No networks auto-detected, using default: 192.168.1.0/24"
    }
    
    return $networks
}

function Show-NetworkSelectionMenu {
    Write-Host "`nWOL Manager - Selectare Retea" -ForegroundColor Yellow
    Write-Host "===============================" -ForegroundColor Yellow
    
    # Try to detect local networks
    $detectedNetworks = @()
    $adapters = Get-NetIPAddress -AddressFamily IPv4 | Where-Object { 
        $_.IPAddress -ne "127.0.0.1" -and ($_.PrefixOrigin -eq "Dhcp" -or $_.PrefixOrigin -eq "Manual")
    }
    
    foreach ($adapter in $adapters) {
        $parts = $adapter.IPAddress.Split('.')
        if ($parts.Length -eq 4) {
            $prefixLength = $adapter.PrefixLength
            
            if ($prefixLength -eq 24) {
                $networkAddr = "$($parts[0]).$($parts[1]).$($parts[2]).0/24"
            } elseif ($prefixLength -eq 20) {
                $thirdOctet = [int]$parts[2]
                $networkThird = $thirdOctet - ($thirdOctet % 16)
                $networkAddr = "$($parts[0]).$($parts[1]).$networkThird.0/20"
            } elseif ($prefixLength -eq 16) {
                $networkAddr = "$($parts[0]).$($parts[1]).0.0/16"
            } else {
                # Fallback to /24
                $networkAddr = "$($parts[0]).$($parts[1]).$($parts[2]).0/24"
            }
            
            if ($detectedNetworks -notcontains $networkAddr) {
                $detectedNetworks += $networkAddr
            }
        }
    }
    
    # Show menu
    Write-Host "`nSelecteaza reteaua de scanat:" -ForegroundColor White
    Write-Host ""
    
    $menuOptions = @()
    $optionIndex = 1
    
    # Add detected networks
    if ($detectedNetworks.Count -gt 0) {
        Write-Host "Retele detectate automat:" -ForegroundColor Green
        foreach ($detectedNetwork in $detectedNetworks) {
            Write-Host "  $optionIndex. $detectedNetwork" -ForegroundColor Cyan
            $menuOptions += $detectedNetwork
            $optionIndex++
        }
        Write-Host ""
    }
    
    # Add common networks
    Write-Host "Retele comune:" -ForegroundColor Green
    $commonNetworks = @("192.168.1.0/24", "192.168.0.0/24", "10.0.0.0/24", "10.0.20.0/24", "172.16.0.0/24")
    
    foreach ($commonNetwork in $commonNetworks) {
        if ($detectedNetworks -notcontains $commonNetwork) {
            Write-Host "  $optionIndex. $commonNetwork" -ForegroundColor Cyan
            $menuOptions += $commonNetwork
            $optionIndex++
        }
    }
    
    Write-Host ""
    Write-Host "  $optionIndex. Introdu manual reteaua" -ForegroundColor Yellow
    $menuOptions += "custom"
    
    Write-Host ""
    Write-Host "  0. Iesire" -ForegroundColor Red
    Write-Host ""
    
    # Get user choice
    do {
        $choice = Read-Host "Selecteaza optiunea (0-$($menuOptions.Count))"
        
        # Convert to integer for comparison
        try {
            $choiceNum = [int]$choice
        } catch {
            Write-Host "Te rog introdu un numar valid!" -ForegroundColor Red
            continue
        }
        
        if ($choiceNum -eq 0) {
            Write-Host "Scanare anulata." -ForegroundColor Yellow
            exit 0
        }
        elseif ($choiceNum -ge 1 -and $choiceNum -le $menuOptions.Count) {
            if ($menuOptions[$choiceNum - 1] -eq "custom") {
                do {
                    $customNetwork = Read-Host "Introdu reteaua in format CIDR (ex: 192.168.100.0/24)"
                    if ($customNetwork -match '^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}/\d{1,2}$') {
                        return @($customNetwork)
                    } else {
                        Write-Host "Format invalid! Foloseste formatul: 192.168.1.0/24" -ForegroundColor Red
                    }
                } while ($true)
            } else {
                return @($menuOptions[$choiceNum - 1])
            }
        } else {
            Write-Host "Optiune invalida! Selecteaza intre 0 si $($menuOptions.Count)." -ForegroundColor Red
        }
    } while ($true)
}

function Test-NetworkAddress {
    param(
        [string]$IPAddress,
        [int]$TimeoutMs = 1000
    )
    
    try {
        $ping = New-Object System.Net.NetworkInformation.Ping
        $result = $ping.Send($IPAddress, $TimeoutMs)
        return $result.Status -eq "Success"
    }
    catch {
        return $false
    }
}

function Get-LocalHostMacAddress {
    param([string]$IPAddress)
    
    Write-Log "Attempting to find MAC address for potential local host: $IPAddress"
    
    try {
        # Method 1: Check if this IP matches any of the local network adapters
        $localAdapters = Get-NetIPAddress -AddressFamily IPv4 | Where-Object { 
            $_.IPAddress -ne "127.0.0.1" -and ($_.PrefixOrigin -eq "Dhcp" -or $_.PrefixOrigin -eq "Manual")
        }
        
        foreach ($adapter in $localAdapters) {
            if ($adapter.IPAddress -eq $IPAddress) {
                Write-Log "Found matching local IP address: $IPAddress on interface $($adapter.InterfaceIndex)"
                
                # Get the corresponding network adapter
                $netAdapter = Get-NetAdapter -InterfaceIndex $adapter.InterfaceIndex -ErrorAction SilentlyContinue
                if ($netAdapter -and $netAdapter.MacAddress) {
                    $cleanMac = $netAdapter.MacAddress.Replace('-', ':').ToLower()
                    Write-Log "Found local MAC address: $cleanMac for IP $IPAddress"
                    return $cleanMac
                }
            }
        }
        
        # Method 2: Use WMI to find network adapters with the matching IP
        $wmiAdapters = Get-WmiObject -Class Win32_NetworkAdapterConfiguration | Where-Object { 
            $_.IPEnabled -eq $true -and $_.IPAddress
        }
        
        foreach ($wmiAdapter in $wmiAdapters) {
            if ($wmiAdapter.IPAddress -contains $IPAddress) {
                if ($wmiAdapter.MACAddress) {
                    $cleanMac = $wmiAdapter.MACAddress.Replace('-', ':').ToLower()
                    Write-Log "Found local MAC via WMI: $cleanMac for IP $IPAddress"
                    return $cleanMac
                }
            }
        }
        
        # Method 3: Check if IP is in the same subnet and try to find the default adapter
        $localIPs = $localAdapters | ForEach-Object { $_.IPAddress }
        foreach ($localIP in $localIPs) {
            $localParts = $localIP -split '\.'
            $targetParts = $IPAddress -split '\.'
            
            # Simple check for same /24 network
            if ($localParts.Length -eq 4 -and $targetParts.Length -eq 4) {
                if ($localParts[0] -eq $targetParts[0] -and 
                    $localParts[1] -eq $targetParts[1] -and 
                    $localParts[2] -eq $targetParts[2]) {
                    
                    # This IP is in the same subnet, find the adapter for the local IP
                    $matchingAdapter = $localAdapters | Where-Object { $_.IPAddress -eq $localIP }
                    if ($matchingAdapter) {
                        $netAdapter = Get-NetAdapter -InterfaceIndex $matchingAdapter.InterfaceIndex -ErrorAction SilentlyContinue
                        if ($netAdapter -and $netAdapter.MacAddress) {
                            $cleanMac = $netAdapter.MacAddress.Replace('-', ':').ToLower()
                            Write-Log "Found MAC for same subnet adapter: $cleanMac (local IP: $localIP, target: $IPAddress)"
                            return $cleanMac
                        }
                    }
                }
            }
        }
        
        Write-Log "Could not determine MAC address for IP: $IPAddress"
        return ""
    }
    catch {
        Write-Log "Error finding local MAC address: $($_.Exception.Message)"
        return ""
    }
}

function Get-NetworkDevices {
    param([string]$NetworkCIDR)
    
    Write-Log "Scanning network: $NetworkCIDR"
    
    # Parse network CIDR
    $parts = $NetworkCIDR -split '/'
    $networkAddress = $parts[0]
    $prefixLength = [int]$parts[1]
    
    # Calculate IP range
    $networkParts = $networkAddress -split '\.'
    $baseIP = "$($networkParts[0]).$($networkParts[1]).$($networkParts[2])"
    
    # Determine scan range based on prefix length
    if ($prefixLength -eq 24) {
        $startIP = 1
        $endIP = 254
    } elseif ($prefixLength -eq 20) {
        # For /20 networks (like 10.0.20.0/20), scan the specific /24 subnet where the network is
        $thirdOctet = [int]$networkParts[2]
        $baseIP = "$($networkParts[0]).$($networkParts[1]).$thirdOctet"
        $startIP = 1
        $endIP = 254
        Write-Log "Scanning /20 network as /24 subnet: $baseIP.0/24"
    } elseif ($prefixLength -eq 16) {
        # For /16 networks, we need to scan multiple /24 subnets
        # For now, just scan the specific /24 subnet where the network address is
        $thirdOctet = [int]$networkParts[2]
        $baseIP = "$($networkParts[0]).$($networkParts[1]).$thirdOctet"
        $startIP = 1
        $endIP = 254
        Write-Log "Scanning /16 network as /24 subnet: $baseIP.0/24"
    } else {
        Write-Log "Unsupported prefix length /$prefixLength. Supported: /16, /20, /24"
        return @()
    }
    
    $devices = @()
    $totalIPs = $endIP - $startIP + 1
    
    # STEP 1: Get ARP table first to find devices that might not respond to ping
    Write-Log "Step 1: Retrieving devices from ARP table..."
    $arpEntries = @{}
    $arpDevices = @{}
    
    try {
        # Use Get-NetNeighbor for Windows 8/Server 2012 and later
        $neighbors = Get-NetNeighbor -AddressFamily IPv4 | Where-Object { $_.State -ne "Unreachable" }
        foreach ($neighbor in $neighbors) {
            if ($neighbor.LinkLayerAddress -and $neighbor.LinkLayerAddress -ne "00-00-00-00-00-00") {
                $cleanMac = $neighbor.LinkLayerAddress.Replace('-', ':').ToLower()
                $arpEntries[$neighbor.IPAddress] = $cleanMac
                
                # Check if this IP is in our scan range
                $ipParts = $neighbor.IPAddress -split '\.'
                if ($ipParts.Length -eq 4) {
                    $currentBaseIP = "$($ipParts[0]).$($ipParts[1]).$($ipParts[2])"
                    $lastOctet = [int]$ipParts[3]
                    
                    if ($currentBaseIP -eq $baseIP -and $lastOctet -ge $startIP -and $lastOctet -le $endIP) {
                        $arpDevices[$neighbor.IPAddress] = @{
                            ip = $neighbor.IPAddress
                            mac = $cleanMac
                            hostname = ""
                            status = "arp"
                        }
                        Write-Log "Found in ARP: $($neighbor.IPAddress) -> $cleanMac"
                    }
                }
            }
        }
    }
    catch {
        Write-Log "Get-NetNeighbor failed, trying arp command..."
        
        # Fallback to arp command
        try {
            $arpOutput = arp -a
            foreach ($line in $arpOutput) {
                if ($line -match '(\d+\.\d+\.\d+\.\d+)\s+([0-9a-fA-F-]{17})\s+\w+') {
                    $ip = $matches[1]
                    $mac = $matches[2].Replace('-', ':').ToLower()
                    if ($mac -ne "00:00:00:00:00:00") {
                        $arpEntries[$ip] = $mac
                        
                        # Check if this IP is in our scan range
                        $ipParts = $ip -split '\.'
                        if ($ipParts.Length -eq 4) {
                            $currentBaseIP = "$($ipParts[0]).$($ipParts[1]).$($ipParts[2])"
                            $lastOctet = [int]$ipParts[3]
                            
                            if ($currentBaseIP -eq $baseIP -and $lastOctet -ge $startIP -and $lastOctet -le $endIP) {
                                $arpDevices[$ip] = @{
                                    ip = $ip
                                    mac = $mac
                                    hostname = ""
                                    status = "arp"
                                }
                                Write-Log "Found in ARP: $ip -> $mac"
                            }
                        }
                    }
                }
            }
        }
        catch {
            Write-Log "Warning: Could not retrieve ARP table"
        }
    }
    
    Write-Log "Found $($arpDevices.Count) devices in ARP table for this network"
    
    # STEP 2: Ping sweep for remaining IPs not found in ARP
    Write-Log "Step 2: Starting ping sweep for remaining addresses..."
    
    $aliveIPs = @()
    $processed = 0
    $skippedCount = 0
    
    # Process IPs in smaller batches to avoid overwhelming the system
    for ($batch = $startIP; $batch -le $endIP; $batch += $BatchSize) {
        $batchEnd = [Math]::Min($batch + $BatchSize - 1, $endIP)
        $jobs = @()
        
        Write-Log "Batch progress: $([Math]::Floor(($batch - $startIP) * 100 / $totalIPs))% - Scanning IPs $batch to $batchEnd"
        
        # Create jobs for current batch (skip IPs already found in ARP)
        for ($i = $batch; $i -le $batchEnd; $i++) {
            $ip = "$baseIP.$i"
            
            if ($arpDevices.ContainsKey($ip)) {
                $processed++
                $skippedCount++
                continue  # Skip IPs already found in ARP
            }
            
            $job = Start-Job -ScriptBlock {
                param($IPAddress, $TimeoutMs)
                try {
                    $ping = New-Object System.Net.NetworkInformation.Ping
                    $result = $ping.Send($IPAddress, $TimeoutMs)
                    return @{
                        IP = $IPAddress
                        Success = ($result.Status -eq "Success")
                        ResponseTime = if ($result.Status -eq "Success") { $result.RoundtripTime } else { $null }
                    }
                }
                catch {
                    return @{
                        IP = $IPAddress
                        Success = $false
                        ResponseTime = $null
                    }
                }
            } -ArgumentList $ip, $TimeoutMs
            
            $jobs += @{ Job = $job; IP = $ip }
        }
        
        # Wait for batch to complete with timeout
        $maxWaitTime = ($TimeoutMs * 2) / 1000  # Convert to seconds, double the ping timeout
        foreach ($jobInfo in $jobs) {
            $result = Wait-Job -Job $jobInfo.Job -Timeout $maxWaitTime | Receive-Job
            Remove-Job -Job $jobInfo.Job -Force
            
            if ($result -and $result.Success) {
                $aliveIPs += $result.IP
                Write-Log "Found alive host via ping: $($result.IP) ($($result.ResponseTime)ms)"
            }
            
            $processed++
        }
        
        # Show progress more frequently
        $progressPercent = [Math]::Floor($processed * 100 / $totalIPs)
        Write-Log "Progress: $processed/$totalIPs addresses scanned ($progressPercent%) - Skipped $skippedCount (already in ARP)"
        
        # Small delay between batches to prevent system overload
        Start-Sleep -Milliseconds 100
    }
    
    Write-Log "Found $($aliveIPs.Count) additional hosts via ping"
    
    # STEP 3: Combine results and resolve hostnames
    Write-Log "Step 3: Building final device list..."
    
    # Add devices found in ARP
    foreach ($arpDevice in $arpDevices.Values) {
        $arpDevice.status = "online"  # Change from "arp" to "online"
        $devices += $arpDevice
    }
    
    # Add devices found via ping (and try to get their MAC from ARP if available)
    foreach ($ip in $aliveIPs) {
        $device = @{
            ip = $ip
            mac = if ($arpEntries.ContainsKey($ip)) { $arpEntries[$ip] } else { "" }
            hostname = ""
            status = "online"
        }
        $devices += $device
    }
    
    # Resolve hostnames and fix MAC addresses for local host
    foreach ($device in $devices) {
        try {
            $hostname = [System.Net.Dns]::GetHostEntry($device.ip).HostName
            if ($hostname -and $hostname -ne $device.ip) {
                $device.hostname = $hostname
            }
        }
        catch {
            # Hostname resolution failed, leave empty
        }
        
        # Special handling for local host (Windows machine running the script)
        if (-not $device.mac -or $device.mac -eq "") {
            $device.mac = Get-LocalHostMacAddress -IPAddress $device.ip
            if ($device.mac) {
                Write-Log "Found local host MAC: $($device.ip) -> $($device.mac)"
            }
        }
        
        if ($device.mac) {
            Write-Log "Final device: $($device.ip) -> $($device.mac) ($($device.hostname))"
        } else {
            Write-Log "Final device: $($device.ip) (no MAC address available) ($($device.hostname))"
        }
    }
    
    return $devices
}

# Main execution
try {
    Write-Log "Starting Windows network scan..."
    
    # Determine networks to scan
    $networksToScan = @()
    
    if ($Network) {
        # Use specified network
        $networksToScan += $Network
    } else {
        # Show interactive menu for network selection
        $networksToScan = Show-NetworkSelectionMenu
    }
    
    if ($networksToScan.Count -eq 0) {
        Write-Error "No networks to scan found"
        exit 1
    }
    
    $allDevices = @()
    
    # Scan each network
    foreach ($net in $networksToScan) {
        Write-Log "Processing network: $net"
        $devices = Get-NetworkDevices -NetworkCIDR $net
        $allDevices += $devices
    }
    
    # Prepare result object
    $result = @{
        success = $true
        timestamp = (Get-Date -Format "yyyy-MM-ddTHH:mm:ss.fffZ")
        networks_scanned = $networksToScan
        computers = $allDevices
        message = "Scanare completata cu succes. Gasite $($allDevices.Count) dispozitive."
    }
    
    # Ensure output directory exists
    $outputDir = Split-Path -Parent $OutputPath
    if (-not (Test-Path $outputDir)) {
        New-Item -ItemType Directory -Path $outputDir -Force | Out-Null
    }
    
    # Save results to JSON
    $result | ConvertTo-Json -Depth 10 | Out-File -FilePath $OutputPath -Encoding UTF8
    
    Write-Log "Results saved to: $OutputPath"
    Write-Log "Scan completed successfully. Found $($allDevices.Count) devices."
    
    # Also output to console for immediate feedback
    if ($allDevices.Count -gt 0) {
        Write-Host "`nDevices found:" -ForegroundColor Yellow
        foreach ($device in $allDevices) {
            $macInfo = if ($device.mac) { $device.mac } else { "No MAC" }
            $hostInfo = if ($device.hostname) { " ($($device.hostname))" } else { "" }
            Write-Host "  $($device.ip) -> $macInfo$hostInfo" -ForegroundColor Cyan
        }
    }
}
catch {
    $errorResult = @{
        success = $false
        timestamp = (Get-Date -Format "yyyy-MM-ddTHH:mm:ss.fffZ")
        message = "Eroare la scanare: $($_.Exception.Message)"
        computers = @()
    }
    
    $errorResult | ConvertTo-Json -Depth 10 | Out-File -FilePath $OutputPath -Encoding UTF8
    
    Write-Error "Scan failed: $($_.Exception.Message)"
    exit 1
}