. (Join-Path $PSScriptRoot 'config.ps1')
. (Join-Path $PSScriptRoot 'logger.ps1')
. (Join-Path $PSScriptRoot 'proxy-test.ps1')

$Script:ProxySuccessCount = 0
$Script:ProxyFailureCount = 0
$Script:LastSwitchTime = (Get-Date).AddSeconds(-1 * $Global:SwitchCooldownSeconds)

function Test-IsAdministrator {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]$identity
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Request-Administrator {
    Write-Log -Message "Requesting Administrator privileges..."
    Start-Process pwsh -ArgumentList "-NoExit", "-File", $PSCommandPath -Verb RunAs
    exit
}

function Get-TargetAdapter {
    $adapter = Get-NetAdapter |
        Where-Object { $_.Name.Contains($Global:TargetNetworkAdapterKeyword) -and $_.Status -eq "Up" } |
        Select-Object -First 1

    if (-not $adapter) {
        throw "No active network adapter found by keyword '$($Global:TargetNetworkAdapterKeyword)'."
    }

    $Global:TargetNetworkAdapter = $adapter
    Write-Log -Message "Target adapter: $($adapter.Name), InterfaceIndex: $($adapter.ifIndex)"
    return $adapter
}

function Get-CurrentDefaultGateway {
    param(
        [Parameter(Mandatory)]
        [int]$InterfaceIndex
    )

    $route = Get-NetRoute -InterfaceIndex $InterfaceIndex -DestinationPrefix "0.0.0.0/0" -ErrorAction SilentlyContinue |
        Sort-Object RouteMetric, ifMetric |
        Select-Object -First 1

    if ($route) {
        return $route.NextHop
    }

    return $null
}

function Get-CurrentMode {
    param(
        [Parameter(Mandatory)]
        [int]$InterfaceIndex
    )

    $gateway = Get-CurrentDefaultGateway -InterfaceIndex $InterfaceIndex
    if ($gateway -eq $Global:CustomGateway) {
        return "PROXY"
    }

    return "DIRECT"
}

function Test-CooldownExpired {
    $elapsed = ((Get-Date) - $Script:LastSwitchTime).TotalSeconds
    return $elapsed -ge $Global:SwitchCooldownSeconds
}

function Test-HostReachable {
    param(
        [Parameter(Mandatory)]
        [string]$Host
    )

    try {
        return [bool](Test-Connection -ComputerName $Host -Count 1 -Quiet -ErrorAction Stop)
    } catch {
        Write-Log -Message "Host check failed for ${Host}: $_"
        return $false
    }
}

function Test-ProxyDns {
    try {
        $result = Resolve-DnsName -Name $Global:TestDomain -Server $Global:CustomGateway -ErrorAction Stop
        return [bool]$result
    } catch {
        Write-Log -Message "Proxy DNS check failed via $($Global:CustomGateway): $_"
        return $false
    }
}

function Test-InternetTcp {
    try {
        return [bool](Test-NetConnection -ComputerName $Global:TestDomain -Port 443 -InformationLevel Quiet)
    } catch {
        Write-Log -Message "Internet TCP check failed for $($Global:TestDomain):443: $_"
        return $false
    }
}

function Test-ProxyHealthy {
    $gatewayReachable = Test-HostReachable -Host $Global:CustomGateway
    if (-not $gatewayReachable) {
        Write-Log -Message "Proxy gateway $($Global:CustomGateway) is not reachable."
        return $false
    }

    $dnsOk = Test-ProxyDns
    if (-not $dnsOk) {
        return $false
    }

    if ($Global:UseSocksProxyHealthCheck) {
        return Test-SocksProxy -ProxyHost $Global:CustomGateway -ProxyPort $Global:ProxyPort -TestHost $Global:TestDomain -TestPort 443
    }

    return $true
}

function Set-DefaultGateway {
    param(
        [Parameter(Mandatory)]
        [int]$InterfaceIndex,

        [Parameter(Mandatory)]
        [string]$Gateway
    )

    $currentGateway = Get-CurrentDefaultGateway -InterfaceIndex $InterfaceIndex
    if ($currentGateway -eq $Gateway) {
        Write-Log -Message "Default gateway is already $Gateway."
        return
    }

    $existingRoutes = Get-NetRoute -InterfaceIndex $InterfaceIndex -DestinationPrefix "0.0.0.0/0" -ErrorAction SilentlyContinue
    foreach ($route in $existingRoutes) {
        Write-Log -Message "Removing default route $($route.NextHop) from interface $InterfaceIndex."
        Remove-NetRoute -InterfaceIndex $InterfaceIndex -DestinationPrefix "0.0.0.0/0" -NextHop $route.NextHop -Confirm:$false
    }

    Write-Log -Message "Adding default route $Gateway to interface $InterfaceIndex."
    New-NetRoute -InterfaceIndex $InterfaceIndex -DestinationPrefix "0.0.0.0/0" -NextHop $Gateway -RouteMetric 1 | Out-Null
}

function Set-ProxyDns {
    param(
        [Parameter(Mandatory)]
        [int]$InterfaceIndex
    )

    Write-Log -Message "Setting DNS to proxy gateway $($Global:CustomGateway)."
    Set-DnsClientServerAddress -InterfaceIndex $InterfaceIndex -ServerAddresses $Global:CustomGateway
    Clear-DnsClientCache
}

function Set-DirectDns {
    param(
        [Parameter(Mandatory)]
        [int]$InterfaceIndex
    )

    if ($Global:DirectDnsServers -and $Global:DirectDnsServers.Count -gt 0) {
        Write-Log -Message "Setting DNS to $($Global:DirectDnsServers -join ', ')."
        Set-DnsClientServerAddress -InterfaceIndex $InterfaceIndex -ServerAddresses $Global:DirectDnsServers
    } else {
        Write-Log -Message "Resetting DNS to DHCP provided servers."
        Set-DnsClientServerAddress -InterfaceIndex $InterfaceIndex -ResetServerAddresses
    }

    Clear-DnsClientCache
}

function Switch-ToProxy {
    param(
        [Parameter(Mandatory)]
        [int]$InterfaceIndex
    )

    Write-Log -Message "Switching to proxy gateway $($Global:CustomGateway)." -ToEventLog
    Set-ProxyDns -InterfaceIndex $InterfaceIndex
    Set-DefaultGateway -InterfaceIndex $InterfaceIndex -Gateway $Global:CustomGateway
    $Global:isCustomGateway = $true
    $Script:LastSwitchTime = Get-Date
}

function Switch-ToDirect {
    param(
        [Parameter(Mandatory)]
        [int]$InterfaceIndex
    )

    Write-Log -Message "Switching to direct gateway $($Global:originalGateway)." -ToEventLog
    Set-DefaultGateway -InterfaceIndex $InterfaceIndex -Gateway $Global:originalGateway
    Set-DirectDns -InterfaceIndex $InterfaceIndex
    $Global:isCustomGateway = $false
    $Script:LastSwitchTime = Get-Date
}

function Initialize-GatewayKeepAlive {
    if (-not (Test-IsAdministrator)) {
        Request-Administrator
    }

    Write-Log -Message "GateWay-KeepAliveForPowershell is start" -ToEventLog
    Write-Log -Message "Running with Administrator privileges."

    $adapter = Get-TargetAdapter
    $mode = Get-CurrentMode -InterfaceIndex $adapter.ifIndex
    $Global:isCustomGateway = $mode -eq "PROXY"
    Write-Log -Message "Initial mode: $mode"

    return $adapter
}

function Start-GatewayKeepAliveLoop {
    param(
        [Parameter(Mandatory)]
        $Adapter
    )

    while ($true) {
        $mode = Get-CurrentMode -InterfaceIndex $Adapter.ifIndex
        $Global:isCustomGateway = $mode -eq "PROXY"
        $proxyHealthy = Test-ProxyHealthy

        if ($mode -eq "PROXY") {
            $internetOk = Test-InternetTcp
            if ($proxyHealthy -and $internetOk) {
                $Script:ProxyFailureCount = 0
            } else {
                $Script:ProxyFailureCount++
                Write-Log -Message "Proxy mode health failed $($Script:ProxyFailureCount)/$($Global:ProxyFailureThreshold)."
            }

            if ($Script:ProxyFailureCount -ge $Global:ProxyFailureThreshold -and (Test-CooldownExpired)) {
                Switch-ToDirect -InterfaceIndex $Adapter.ifIndex
                $Script:ProxyFailureCount = 0
                $Script:ProxySuccessCount = 0
            }
        } else {
            if ($proxyHealthy) {
                $Script:ProxySuccessCount++
                Write-Log -Message "Proxy mode candidate succeeded $($Script:ProxySuccessCount)/$($Global:ProxySuccessThreshold)."
            } else {
                $Script:ProxySuccessCount = 0
            }

            if ($Script:ProxySuccessCount -ge $Global:ProxySuccessThreshold -and (Test-CooldownExpired)) {
                Switch-ToProxy -InterfaceIndex $Adapter.ifIndex
                $Script:ProxySuccessCount = 0
                $Script:ProxyFailureCount = 0
            }
        }

        Start-Sleep -Seconds $Global:HealthCheckIntervalSeconds
    }
}

$adapter = Initialize-GatewayKeepAlive
Start-GatewayKeepAliveLoop -Adapter $adapter