Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

function Assert-Administrator {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]::new($identity)
    if (-not $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
        throw 'Run this script from an elevated PowerShell session.'
    }
}

function New-ActivityWatchDirectory {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path
    )

    if (-not (Test-Path -LiteralPath $Path)) {
        New-Item -Path $Path -ItemType Directory -Force | Out-Null
    }
}

function Get-ActivityWatchPackageUrl {
    param(
        [string]$Version = 'v0.13.2'
    )

    return "https://github.com/ActivityWatch/activitywatch/releases/download/$Version/activitywatch-$Version-windows-x86_64.zip"
}

function Get-ActivityWatchArchive {
    param(
        [string]$PackageZipPath,
        [string]$PackageUrl,
        [string]$Version = 'v0.13.2',
        [Parameter(Mandatory = $true)]
        [string]$WorkingRoot
    )

    New-ActivityWatchDirectory -Path $WorkingRoot

    if ($PackageZipPath) {
        $resolved = Resolve-Path -LiteralPath $PackageZipPath -ErrorAction Stop
        return $resolved.Path
    }

    if (-not $PackageUrl) {
        $PackageUrl = Get-ActivityWatchPackageUrl -Version $Version
    }

    [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12
    $archivePath = Join-Path $WorkingRoot ("activitywatch-{0}.zip" -f $Version.TrimStart('v'))
    Invoke-WebRequest -Uri $PackageUrl -OutFile $archivePath
    return $archivePath
}

function Get-ActivityWatchPackageRoot {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ExpandedRoot
    )

    $afkBinary = Get-ChildItem -Path $ExpandedRoot -Filter 'aw-watcher-afk.exe' -File -Recurse |
        Select-Object -First 1

    if (-not $afkBinary) {
        throw "Cannot find aw-watcher-afk.exe under $ExpandedRoot."
    }

    return (Split-Path -Path (Split-Path -Path $afkBinary.FullName -Parent) -Parent)
}

function Install-ActivityWatchPackage {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ArchivePath,
        [Parameter(Mandatory = $true)]
        [string]$InstallRoot,
        [Parameter(Mandatory = $true)]
        [string]$WorkingRoot,
        [Parameter(Mandatory = $true)]
        [string]$BackupRoot
    )

    New-ActivityWatchDirectory -Path $WorkingRoot
    New-ActivityWatchDirectory -Path $BackupRoot

    $extractRoot = Join-Path $WorkingRoot ('extract-' + [guid]::NewGuid().Guid)
    if (Test-Path -LiteralPath $extractRoot) {
        Remove-Item -LiteralPath $extractRoot -Recurse -Force
    }
    New-ActivityWatchDirectory -Path $extractRoot

    Expand-Archive -Path $ArchivePath -DestinationPath $extractRoot -Force
    $packageRoot = Get-ActivityWatchPackageRoot -ExpandedRoot $extractRoot

    if (Test-Path -LiteralPath $InstallRoot) {
        $existingItems = Get-ChildItem -LiteralPath $InstallRoot -Force -ErrorAction SilentlyContinue
        if ($existingItems) {
            $stamp = Get-Date -Format 'yyyyMMdd-HHmmss'
            $backupPath = Join-Path $BackupRoot ("install-$stamp")
            New-ActivityWatchDirectory -Path $backupPath
            Copy-Item -Path (Join-Path $InstallRoot '*') -Destination $backupPath -Recurse -Force
            Get-ChildItem -LiteralPath $InstallRoot -Force | Remove-Item -Recurse -Force
        }
    }
    else {
        New-ActivityWatchDirectory -Path $InstallRoot
    }

    Copy-Item -Path (Join-Path $packageRoot '*') -Destination $InstallRoot -Recurse -Force

    return [pscustomobject]@{
        PackageRoot = $packageRoot
        ExtractRoot = $extractRoot
        BackupRoot  = $BackupRoot
    }
}

function Get-ActivityWatchExecutableMap {
    param(
        [Parameter(Mandatory = $true)]
        [string]$InstallRoot
    )

    $map = [ordered]@{
        Afk    = Join-Path $InstallRoot 'aw-watcher-afk\aw-watcher-afk.exe'
        Window = Join-Path $InstallRoot 'aw-watcher-window\aw-watcher-window.exe'
    }

    foreach ($entry in $map.GetEnumerator()) {
        if (-not (Test-Path -LiteralPath $entry.Value)) {
            throw "Missing required ActivityWatch binary: $($entry.Value)"
        }
    }

    return [pscustomobject]$map
}

function Normalize-ActivityWatchUsers {
    param(
        [string[]]$Users,
        [string]$UserListPath,
        [string]$Domain
    )

    $collected = New-Object System.Collections.Generic.List[string]

    if ($Users) {
        foreach ($user in $Users) {
            if (-not [string]::IsNullOrWhiteSpace($user)) {
                $collected.Add($user.Trim())
            }
        }
    }

    if ($UserListPath) {
        $resolved = Resolve-Path -LiteralPath $UserListPath -ErrorAction Stop
        $extension = [IO.Path]::GetExtension($resolved.Path)
        if ($extension -ieq '.csv') {
            $rows = Import-Csv -LiteralPath $resolved.Path
            foreach ($row in $rows) {
                foreach ($column in 'User', 'Username', 'SamAccountName', 'Login') {
                    if ($row.PSObject.Properties.Name -contains $column) {
                        $value = [string]$row.$column
                        if (-not [string]::IsNullOrWhiteSpace($value)) {
                            $collected.Add($value.Trim())
                            break
                        }
                    }
                }
            }
        }
        else {
            Get-Content -LiteralPath $resolved.Path | ForEach-Object {
                $line = $_.Trim()
                if ($line -and -not $line.StartsWith('#')) {
                    $collected.Add($line)
                }
            }
        }
    }

    $normalized = $collected |
        Where-Object { -not [string]::IsNullOrWhiteSpace($_) } |
        ForEach-Object {
            if ($Domain -and ($_ -notmatch '[\\@]')) {
                '{0}\{1}' -f $Domain, $_
            }
            else {
                $_
            }
        } |
        Sort-Object -Unique

    if (-not $normalized -or $normalized.Count -eq 0) {
        throw 'No target users resolved. Provide -Users or -UserListPath.'
    }

    return @($normalized)
}

function Get-ActivityWatchTaskNameToken {
    param(
        [Parameter(Mandatory = $true)]
        [string]$UserId
    )

    $buffer = [Text.StringBuilder]::new()
    foreach ($character in $UserId.ToCharArray()) {
        if ([char]::IsLetterOrDigit($character)) {
            [void]$buffer.Append($character)
        }
        else {
            [void]$buffer.Append('_')
        }
    }

    return $buffer.ToString().Trim('_')
}

function New-ActivityWatchUserTaskDefinitions {
    param(
        [Parameter(Mandatory = $true)]
        [string[]]$Users
    )

    $result = foreach ($user in $Users) {
        $token = Get-ActivityWatchTaskNameToken -UserId $user
        [pscustomobject]@{
            UserId         = $user
            LaunchTaskName = "ActivityWatch Launch [$token]"
        }
    }

    return @($result)
}

function Copy-ActivityWatchCollectorAssets {
    param(
        [Parameter(Mandatory = $true)]
        [string]$CollectorScriptSource,
        [Parameter(Mandatory = $true)]
        [string]$EndpointCollectorScriptSource,
        [Parameter(Mandatory = $true)]
        [string]$ExampleRulesSource,
        [Parameter(Mandatory = $true)]
        [string]$ExamplePolicySource,
        [Parameter(Mandatory = $true)]
        [string]$StateRoot,
        [string]$CustomRulesSource,
        [string]$CustomPolicySource
    )

    New-ActivityWatchDirectory -Path $StateRoot

    $collectorTarget = Join-Path $StateRoot 'browser-domains-native-collector.ps1'
    $endpointCollectorTarget = Join-Path $StateRoot 'dlp-endpoint-signals-collector.ps1'
    $exampleRulesTarget = Join-Path $StateRoot 'web-category-rules.example.json'
    $rulesTarget = Join-Path $StateRoot 'web-category-rules.json'
    $examplePolicyTarget = Join-Path $StateRoot 'dlp-policy.example.json'
    $policyTarget = Join-Path $StateRoot 'dlp-policy.json'

    Copy-Item -LiteralPath $CollectorScriptSource -Destination $collectorTarget -Force
    Copy-Item -LiteralPath $EndpointCollectorScriptSource -Destination $endpointCollectorTarget -Force
    Copy-Item -LiteralPath $ExampleRulesSource -Destination $exampleRulesTarget -Force
    Copy-Item -LiteralPath $ExamplePolicySource -Destination $examplePolicyTarget -Force

    if ($CustomRulesSource) {
        $resolvedRules = Resolve-Path -LiteralPath $CustomRulesSource -ErrorAction Stop
        Copy-Item -LiteralPath $resolvedRules.Path -Destination $rulesTarget -Force
    }

    if ($CustomPolicySource) {
        $resolvedPolicy = Resolve-Path -LiteralPath $CustomPolicySource -ErrorAction Stop
        Copy-Item -LiteralPath $resolvedPolicy.Path -Destination $policyTarget -Force
    }
    else {
        Copy-Item -LiteralPath $examplePolicyTarget -Destination $policyTarget -Force
    }

    return [pscustomobject]@{
        CollectorScript         = $collectorTarget
        EndpointCollectorScript = $endpointCollectorTarget
        ExampleRules            = $exampleRulesTarget
        ActiveRules             = $rulesTarget
        ExamplePolicy           = $examplePolicyTarget
        ActivePolicy            = $policyTarget
    }
}

function New-ActivityWatchDeploymentConfig {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ServerHost,
        [Parameter(Mandatory = $true)]
        [int]$ServerPort,
        [Parameter(Mandatory = $true)]
        [string]$ServerScheme,
        [Parameter(Mandatory = $true)]
        [string]$InstallRoot,
        [Parameter(Mandatory = $true)]
        [string]$StateRoot,
        [Parameter(Mandatory = $true)]
        [string]$LogsRoot,
        [Parameter(Mandatory = $true)]
        [string]$CollectorScript,
        [Parameter(Mandatory = $true)]
        [string]$EndpointCollectorScript,
        [Parameter(Mandatory = $true)]
        [string]$RulesPath,
        [Parameter(Mandatory = $true)]
        [string]$PolicyPath,
        [Parameter(Mandatory = $true)]
        [int]$PollSeconds,
        [Parameter(Mandatory = $true)]
        [int]$PulseSeconds,
        [Parameter(Mandatory = $true)]
        [int]$RecoveryIntervalSeconds,
        [bool]$AfkEnabled = $true,
        [bool]$WindowEnabled = $true,
        [bool]$LocalAgentLogsEnabled = $true,
        [bool]$IncidentCaptureEnabled = $true,
        [bool]$IncidentScreenshotEnabled = $true,
        [string]$IncidentArtifactsRoot,
        [bool]$LogonMarkerEnabled = $true,
        [Parameter(Mandatory = $true)]
        [string]$LaunchScriptPath,
        [Parameter(Mandatory = $true)]
        [string]$RecoveryScriptPath,
        [Parameter(Mandatory = $true)]
        [pscustomobject[]]$UserTasks,
        [string]$PackageVersion = 'v0.13.2'
    )

    $effectiveIncidentArtifactsRoot = if ($IncidentArtifactsRoot) { $IncidentArtifactsRoot } else { Join-Path $StateRoot 'incident-artifacts' }

    return [pscustomobject]@{
        version  = 1
        generatedAtUtc = (Get-Date).ToUniversalTime().ToString('o')
        server   = [pscustomobject]@{
            host   = $ServerHost
            port   = $ServerPort
            scheme = $ServerScheme
        }
        paths    = [pscustomobject]@{
            installRoot    = $InstallRoot
            stateRoot      = $StateRoot
            logsRoot       = $LogsRoot
            collectorScript = $CollectorScript
            endpointCollectorScript = $EndpointCollectorScript
            rulesPath      = $RulesPath
            policyPath     = $PolicyPath
            launchScript   = $LaunchScriptPath
            recoveryScript = $RecoveryScriptPath
        }
        collector = [pscustomobject]@{
            pollSeconds  = $PollSeconds
            pulseSeconds = $PulseSeconds
        }
        collectors = [pscustomobject]@{
            afkEnabled   = $AfkEnabled
            windowEnabled = $WindowEnabled
        }
        logging = [pscustomobject]@{
            localAgentLogsEnabled = $LocalAgentLogsEnabled
        }
        incidentCapture = [pscustomobject]@{
            enabled           = $IncidentCaptureEnabled
            screenshotEnabled = $IncidentScreenshotEnabled
            artifactsRoot     = $effectiveIncidentArtifactsRoot
        }
        sessionEvents = [pscustomobject]@{
            logonEnabled = $LogonMarkerEnabled
            bucketPrefix = 'aw-session-events'
        }
        recovery = [pscustomobject]@{
            intervalSeconds = $RecoveryIntervalSeconds
            taskName        = 'ActivityWatch Recovery'
        }
        dlp = [pscustomobject]@{
            incidentBucketPrefix = 'aw-dlp-incidents'
            enabled              = $true
        }
        package = [pscustomobject]@{
            version = $PackageVersion
        }
        userTasks = @($UserTasks)
    }
}

function Write-ActivityWatchDeploymentConfig {
    param(
        [Parameter(Mandatory = $true)]
        [pscustomobject]$Config,
        [Parameter(Mandatory = $true)]
        [string]$Path
    )

    $directory = Split-Path -Path $Path -Parent
    if ($directory) {
        New-ActivityWatchDirectory -Path $directory
    }

    $json = $Config | ConvertTo-Json -Depth 8
    Set-Content -LiteralPath $Path -Value $json -Encoding UTF8
}

function Read-ActivityWatchDeploymentConfig {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path
    )

    if (-not (Test-Path -LiteralPath $Path)) {
        throw "Deployment config not found: $Path"
    }

    return Get-Content -LiteralPath $Path -Raw | ConvertFrom-Json
}

function Write-ActivityWatchLaunchScript {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path,
        [Parameter(Mandatory = $true)]
        [string]$ConfigPath
    )

    $content = @"
param(
    [string]`$ConfigPath = '$ConfigPath'
)

Set-StrictMode -Version Latest
`$ErrorActionPreference = 'Stop'

function Get-DeploymentConfig {
    param([string]`$Path)
    return Get-Content -LiteralPath `$Path -Raw | ConvertFrom-Json
}

function Test-ProcessInSession {
    param(
        [string]`$Name,
        [int]`$SessionId
    )

    return [bool](Get-Process -Name `$Name -ErrorAction SilentlyContinue | Where-Object { `$_.SessionId -eq `$SessionId } | Select-Object -First 1)
}

function Test-CollectorRunning {
    param(
        [string]`$ScriptPath,
        [int]`$SessionId
    )

    `$escapedCollector = [Regex]::Escape(`$ScriptPath)
    `$processes = Get-CimInstance Win32_Process -ErrorAction SilentlyContinue |
        Where-Object {
            (`$_.Name -ieq 'powershell.exe' -or `$_.Name -ieq 'pwsh.exe') -and
            `$_.SessionId -eq `$SessionId -and
            `$_.CommandLine -match `$escapedCollector
        }

    return [bool](`$processes | Select-Object -First 1)
}

function Invoke-AwJsonPost {
    param(
        [Parameter(Mandatory = `$true)][string]`$Uri,
        [Parameter(Mandatory = `$true)][string]`$Json
    )

    `$bytes = [Text.Encoding]::UTF8.GetBytes(`$Json)
    Invoke-RestMethod -Method Post -Uri `$Uri -ContentType 'application/json; charset=utf-8' -Body `$bytes | Out-Null
}

function Ensure-Bucket {
    param(
        [string]`$BucketId,
        [string]`$ClientName,
        [string]`$BucketType
    )

    if (`$script:KnownBuckets.ContainsKey(`$BucketId)) {
        return
    }

    try {
        Invoke-RestMethod -Method Get -Uri "`$(`$script:ApiBase)/buckets/`$BucketId" | Out-Null
        `$script:KnownBuckets[`$BucketId] = `$true
        return
    }
    catch {
    }

    `$body = @{
        client   = `$ClientName
        type     = `$BucketType
        hostname = `$script:Hostname
    } | ConvertTo-Json -Compress

    try {
        Invoke-AwJsonPost -Uri "`$(`$script:ApiBase)/buckets/`$BucketId" -Json `$body
    }
    catch {
        try {
            Invoke-RestMethod -Method Get -Uri "`$(`$script:ApiBase)/buckets/`$BucketId" | Out-Null
        }
        catch {
            return
        }
    }

    `$script:KnownBuckets[`$BucketId] = `$true
}

function Send-LogonMarkerIfNeeded {
    param(
        [pscustomobject]`$Config,
        [int]`$SessionId
    )

    `$sessionEvents = if (`$Config.PSObject.Properties.Name -contains 'sessionEvents') { `$Config.sessionEvents } else { `$null }
    `$logging = if (`$Config.PSObject.Properties.Name -contains 'logging') { `$Config.logging } else { `$null }
    `$logonEnabled = if (`$sessionEvents -and `$sessionEvents.PSObject.Properties.Name -contains 'logonEnabled') { [bool]`$sessionEvents.logonEnabled } else { `$false }
    if (-not `$logonEnabled) {
        return
    }

    `$bucketPrefix = if (`$sessionEvents -and `$sessionEvents.PSObject.Properties.Name -contains 'bucketPrefix' -and -not [string]::IsNullOrWhiteSpace([string]`$sessionEvents.bucketPrefix)) {
        [string]`$sessionEvents.bucketPrefix
    }
    else {
        'aw-session-events'
    }

    `$stateRoot = [string]`$Config.paths.stateRoot
    `$markerRoots = New-Object System.Collections.Generic.List[string]
    if (-not [string]::IsNullOrWhiteSpace(`$env:LOCALAPPDATA)) {
        `$markerRoots.Add((Join-Path `$env:LOCALAPPDATA 'ActivityWatch-Phase2\markers'))
    }
    if (-not [string]::IsNullOrWhiteSpace(`$stateRoot)) {
        `$markerRoots.Add((Join-Path `$stateRoot 'markers'))
    }

    `$markerDir = `$null
    foreach (`$candidate in `$markerRoots) {
        try {
            if (-not (Test-Path -LiteralPath `$candidate)) {
                New-Item -Path `$candidate -ItemType Directory -Force | Out-Null
            }

            `$probePath = Join-Path `$candidate 'write-test.tmp'
            Set-Content -LiteralPath `$probePath -Value 'ok' -Encoding ASCII
            Remove-Item -LiteralPath `$probePath -Force -ErrorAction SilentlyContinue
            `$markerDir = `$candidate
            break
        }
        catch {
        }
    }

    if (-not `$markerDir) {
        return
    }

    `$markerFile = Join-Path `$markerDir ("logon-{0}-{1}.marker" -f `$env:USERNAME, `$SessionId)
    if (Test-Path -LiteralPath `$markerFile) {
        return
    }

    Set-Content -LiteralPath `$markerFile -Value ((Get-Date).ToUniversalTime().ToString('o')) -Encoding UTF8

    `$bucketId = ('{0}_{1}' -f `$bucketPrefix, `$script:Hostname)
    Ensure-Bucket -BucketId `$bucketId -ClientName 'aw-session-events' -BucketType 'aw.session.event'

    `$payload = @{
        timestamp = (Get-Date).ToUniversalTime().ToString('yyyy-MM-ddTHH:mm:ss.fffZ')
        duration  = 0
        data      = @{
            eventType = 'logon'
            username  = `$env:USERNAME
            userId    = "`$(`$env:USERDOMAIN)\`$(`$env:USERNAME)"
            sessionId = `$SessionId
            hostname  = `$script:Hostname
            source    = 'launch-watchers-phase2'
        }
    } | ConvertTo-Json -Depth 5 -Compress

    try {
        Invoke-AwJsonPost -Uri "`$(`$script:ApiBase)/buckets/`$bucketId/heartbeat?pulsetime=1" -Json `$payload
    }
    catch {
        Remove-Item -LiteralPath `$markerFile -Force -ErrorAction SilentlyContinue
        throw
    }
}

function Start-CollectorScriptIfNeeded {
    param(
        [string]`$ScriptPath,
        [string]`$ConfigPath,
        [string]`$PowerShellExe,
        [int]`$SessionId
    )

    if (-not (Test-Path -LiteralPath `$ScriptPath)) {
        return
    }

    if (Test-CollectorRunning -ScriptPath `$ScriptPath -SessionId `$SessionId) {
        return
    }

    Start-Process -FilePath `$PowerShellExe -ArgumentList @(
        '-NoProfile',
        '-WindowStyle', 'Hidden',
        '-ExecutionPolicy', 'Bypass',
        '-File', `$ScriptPath,
        '-ConfigPath', `$ConfigPath
    ) -WindowStyle Hidden
}

`$config = Get-DeploymentConfig -Path `$ConfigPath
`$sessionId = (Get-Process -Id `$PID).SessionId
`$installRoot = [string]`$config.paths.installRoot
`$script:ApiBase = '{0}://{1}:{2}/api/0' -f [string]`$config.server.scheme, [string]`$config.server.host, [string]`$config.server.port
`$script:Hostname = `$env:COMPUTERNAME
`$script:KnownBuckets = @{}
`$collectorScript = [string]`$config.paths.collectorScript
`$endpointCollectorScript = if (`$config.paths.PSObject.Properties.Name -contains 'endpointCollectorScript') { [string]`$config.paths.endpointCollectorScript } else { '' }
`$afkExe = Join-Path `$installRoot 'aw-watcher-afk\aw-watcher-afk.exe'
`$windowExe = Join-Path `$installRoot 'aw-watcher-window\aw-watcher-window.exe'
`$serverArgs = @('--host', [string]`$config.server.host, '--port', [string]`$config.server.port)
`$powershellExe = Join-Path `$env:SystemRoot 'System32\WindowsPowerShell\v1.0\powershell.exe'
`$afkEnabled = if (`$config.PSObject.Properties.Name -contains 'collectors' -and `$config.collectors.PSObject.Properties.Name -contains 'afkEnabled') { [bool]`$config.collectors.afkEnabled } else { `$true }
`$windowEnabled = if (`$config.PSObject.Properties.Name -contains 'collectors' -and `$config.collectors.PSObject.Properties.Name -contains 'windowEnabled') { [bool]`$config.collectors.windowEnabled } else { `$true }

if (`$afkEnabled -and -not (Test-Path -LiteralPath `$afkExe)) {
    throw "Missing aw-watcher-afk.exe: `$afkExe"
}

if (`$windowEnabled -and -not (Test-Path -LiteralPath `$windowExe)) {
    throw "Missing aw-watcher-window.exe: `$windowExe"
}

if (`$afkEnabled -and -not (Test-ProcessInSession -Name 'aw-watcher-afk' -SessionId `$sessionId)) {
    Start-Process -FilePath `$afkExe -ArgumentList `$serverArgs -WindowStyle Hidden
}

if (`$windowEnabled -and -not (Test-ProcessInSession -Name 'aw-watcher-window' -SessionId `$sessionId)) {
    Start-Process -FilePath `$windowExe -ArgumentList `$serverArgs -WindowStyle Hidden
}

try {
    Send-LogonMarkerIfNeeded -Config `$config -SessionId `$sessionId
}
catch {
}
Start-CollectorScriptIfNeeded -ScriptPath `$collectorScript -ConfigPath `$ConfigPath -PowerShellExe `$powershellExe -SessionId `$sessionId
Start-CollectorScriptIfNeeded -ScriptPath `$endpointCollectorScript -ConfigPath `$ConfigPath -PowerShellExe `$powershellExe -SessionId `$sessionId
"@

    Set-Content -LiteralPath $Path -Value $content -Encoding UTF8
}

function Write-ActivityWatchRecoveryScript {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path,
        [Parameter(Mandatory = $true)]
        [string]$ConfigPath
    )

    $content = @"
param(
    [string]`$ConfigPath = '$ConfigPath'
)

Set-StrictMode -Version Latest
`$ErrorActionPreference = 'Continue'

function Get-DeploymentConfig {
    param([string]`$Path)
    return Get-Content -LiteralPath `$Path -Raw | ConvertFrom-Json
}

function Get-RecoveryConfigPaths {
    param([string]`$PrimaryConfigPath)

    `$paths = New-Object System.Collections.Generic.List[string]
    if (`$PrimaryConfigPath -and (Test-Path -LiteralPath `$PrimaryConfigPath)) {
        `$paths.Add((Resolve-Path -LiteralPath `$PrimaryConfigPath).Path)
    }

    `$searchRoot = `$env:ProgramData
    if (`$PrimaryConfigPath) {
        `$stateRoot = Split-Path -Path `$PrimaryConfigPath -Parent
        `$candidateRoot = Split-Path -Path `$stateRoot -Parent
        if (`$candidateRoot -and (Test-Path -LiteralPath `$candidateRoot)) {
            `$searchRoot = `$candidateRoot
        }
    }

    if (Test-Path -LiteralPath `$searchRoot) {
        Get-ChildItem -LiteralPath `$searchRoot -Directory -ErrorAction SilentlyContinue |
            Where-Object { `$_.Name -like 'ActivityWatch*' } |
            ForEach-Object {
                `$candidate = Join-Path `$_.FullName 'deployment-config.json'
                if (Test-Path -LiteralPath `$candidate) {
                    `$paths.Add(`$candidate)
                }
            }
    }

    return @(`$paths | Sort-Object -Unique)
}

function Get-RecoveryTaskNames {
    param([string[]]`$ConfigPaths)

    `$taskNames = New-Object 'System.Collections.Generic.HashSet[string]' ([System.StringComparer]::OrdinalIgnoreCase)
    foreach (`$candidatePath in @(`$ConfigPaths)) {
        try {
            `$config = Get-DeploymentConfig -Path `$candidatePath
            foreach (`$task in @(`$config.userTasks)) {
                `$taskName = [string]`$task.launchTaskName
                if (-not [string]::IsNullOrWhiteSpace(`$taskName)) {
                    [void]`$taskNames.Add(`$taskName)
                }
            }
        }
        catch {
        }
    }

    return @(`$taskNames)
}

while (`$true) {
    `$sleepSeconds = 180
    try {
        `$configPaths = Get-RecoveryConfigPaths -PrimaryConfigPath `$ConfigPath
        foreach (`$taskName in Get-RecoveryTaskNames -ConfigPaths `$configPaths) {
            Start-ScheduledTask -TaskName `$taskName -ErrorAction SilentlyContinue
        }

        `$config = Get-DeploymentConfig -Path `$ConfigPath
        if (`$config -and `$config.recovery -and `$config.recovery.intervalSeconds) {
            `$sleepSeconds = [Math]::Max([int]`$config.recovery.intervalSeconds, 30)
        }
    }
    catch {
    }

    Start-Sleep -Seconds `$sleepSeconds
}
"@

    Set-Content -LiteralPath $Path -Value $content -Encoding UTF8
}

function Get-ActivityWatchHiddenLauncherPath {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ScriptPath
    )

    $directory = Split-Path -Path $ScriptPath -Parent
    $baseName = [IO.Path]::GetFileNameWithoutExtension($ScriptPath)
    return Join-Path $directory ("{0}-hidden.vbs" -f $baseName)
}

function Write-ActivityWatchHiddenPowerShellWrapper {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path,
        [Parameter(Mandatory = $true)]
        [string]$ScriptPath,
        [Parameter(Mandatory = $true)]
        [string]$ConfigPath
    )

    $directory = Split-Path -Path $Path -Parent
    if ($directory) {
        New-ActivityWatchDirectory -Path $directory
    }

    $powershellExe = Join-Path $env:SystemRoot 'System32\WindowsPowerShell\v1.0\powershell.exe'
    $escapedPowerShellExe = $powershellExe.Replace('"', '""')
    $escapedScriptPath = $ScriptPath.Replace('"', '""')
    $escapedConfigPath = $ConfigPath.Replace('"', '""')

    $content = @"
Set shell = CreateObject("WScript.Shell")
shell.Run """$escapedPowerShellExe"" -NoProfile -ExecutionPolicy Bypass -File ""$escapedScriptPath"" -ConfigPath ""$escapedConfigPath""", 0, False
"@

    Set-Content -LiteralPath $Path -Value $content -Encoding ASCII
}

function Remove-LegacyActivityWatchEntries {
    $legacyTaskNames = @(
        'ActivityWatch Watchers',
        'ActivityWatch Guard',
        'ActivityWatch Heal'
    )

    foreach ($taskName in $legacyTaskNames) {
        Unregister-ScheduledTask -TaskName $taskName -Confirm:$false -ErrorAction SilentlyContinue
    }

    $runKey = 'HKLM:\Software\Microsoft\Windows\CurrentVersion\Run'
    foreach ($name in 'ActivityWatchAFK', 'ActivityWatchWindow', 'ActivityWatchBrowserCollector') {
        Remove-ItemProperty -Path $runKey -Name $name -ErrorAction SilentlyContinue
    }
}

function Remove-ActivityWatchScheduledTask {
    param(
        [Parameter(Mandatory = $true)]
        [string]$TaskName
    )

    Unregister-ScheduledTask -TaskName $TaskName -Confirm:$false -ErrorAction SilentlyContinue
    & cmd.exe /c "schtasks /Delete /TN `"$TaskName`" /F >nul 2>&1" | Out-Null

    for ($attempt = 0; $attempt -lt 10; $attempt++) {
        $task = Get-ScheduledTask -TaskName $TaskName -ErrorAction SilentlyContinue
        if (-not $task) {
            return
        }

        Start-Sleep -Milliseconds 300
    }
}

function Set-ActivityWatchScheduledTaskAction {
    param(
        [Parameter(Mandatory = $true)]
        [string]$TaskName,
        [Parameter(Mandatory = $true)]
        [string]$Execute,
        [Parameter(Mandatory = $true)]
        [string]$Arguments
    )

    $taskCommand = ('"{0}" {1}' -f $Execute, $Arguments)
    & schtasks.exe /Change /TN $TaskName /TR $taskCommand | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw "schtasks.exe /Change failed for $TaskName"
    }
}

function Get-ActivityWatchScheduledTaskByCommand {
    param(
        [Parameter(Mandatory = $true)]
        [string]$TaskName,
        [string]$CommandMatch
    )

    $task = Get-ScheduledTask -TaskName $TaskName -ErrorAction SilentlyContinue
    if ($task) {
        return $task
    }

    if ([string]::IsNullOrWhiteSpace($CommandMatch)) {
        return $null
    }

    foreach ($candidate in @(Get-ScheduledTask | Where-Object { $_.TaskName -like 'ActivityWatch Launch*' })) {
        foreach ($action in @($candidate.Actions)) {
            if ([string]$action.Arguments -like "*$CommandMatch*") {
                return $candidate
            }
        }
    }

    return $null
}

function Register-ActivityWatchUserTasks {
    param(
        [Parameter(Mandatory = $true)]
        [pscustomobject[]]$TaskDefinitions,
        [Parameter(Mandatory = $true)]
        [string]$LaunchScriptPath,
        [Parameter(Mandatory = $true)]
        [string]$ConfigPath
    )

    $wscriptExe = Join-Path $env:SystemRoot 'System32\wscript.exe'
    $launcherPath = Get-ActivityWatchHiddenLauncherPath -ScriptPath $LaunchScriptPath
    Write-ActivityWatchHiddenPowerShellWrapper -Path $launcherPath -ScriptPath $LaunchScriptPath -ConfigPath $ConfigPath

    foreach ($definition in $TaskDefinitions) {
        $action = New-ScheduledTaskAction -Execute $wscriptExe -Argument "//B //NoLogo `"$launcherPath`""
        $trigger = New-ScheduledTaskTrigger -AtLogOn -User $definition.UserId
        $principal = New-ScheduledTaskPrincipal -UserId $definition.UserId -LogonType Interactive -RunLevel Highest
        $settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -StartWhenAvailable -MultipleInstances IgnoreNew -ExecutionTimeLimit (New-TimeSpan -Hours 0)
        $existingTask = Get-ActivityWatchScheduledTaskByCommand -TaskName $definition.LaunchTaskName -CommandMatch $ConfigPath

        if ($existingTask) {
            Set-ActivityWatchScheduledTaskAction -TaskName $existingTask.TaskName -Execute $wscriptExe -Arguments $action.Arguments
            continue
        }

        Remove-ActivityWatchScheduledTask -TaskName $definition.LaunchTaskName
        Register-ScheduledTask -TaskName $definition.LaunchTaskName -Action $action -Trigger $trigger -Principal $principal -Settings $settings | Out-Null
    }
}

function Register-ActivityWatchRecoveryTask {
    param(
        [Parameter(Mandatory = $true)]
        [string]$TaskName,
        [Parameter(Mandatory = $true)]
        [string]$RecoveryScriptPath,
        [Parameter(Mandatory = $true)]
        [string]$ConfigPath
    )

    Remove-ActivityWatchScheduledTask -TaskName $TaskName

    $wscriptExe = Join-Path $env:SystemRoot 'System32\wscript.exe'
    $launcherPath = Get-ActivityWatchHiddenLauncherPath -ScriptPath $RecoveryScriptPath
    Write-ActivityWatchHiddenPowerShellWrapper -Path $launcherPath -ScriptPath $RecoveryScriptPath -ConfigPath $ConfigPath
    $action = New-ScheduledTaskAction -Execute $wscriptExe -Argument "//B //NoLogo `"$launcherPath`""
    $trigger = New-ScheduledTaskTrigger -AtStartup
    $principal = New-ScheduledTaskPrincipal -UserId 'SYSTEM' -LogonType ServiceAccount -RunLevel Highest
    $settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -StartWhenAvailable -Hidden -MultipleInstances IgnoreNew -ExecutionTimeLimit (New-TimeSpan -Hours 0)

    Register-ScheduledTask -TaskName $TaskName -Action $action -Trigger $trigger -Principal $principal -Settings $settings | Out-Null
}

function Set-ActivityWatchAcl {
    param(
        [Parameter(Mandatory = $true)]
        [string]$InstallRoot,
        [Parameter(Mandatory = $true)]
        [string]$StateRoot,
        [Parameter(Mandatory = $true)]
        [string]$LogsRoot
    )

    foreach ($path in $InstallRoot, $StateRoot, $LogsRoot) {
        New-ActivityWatchDirectory -Path $path
    }

    & icacls $InstallRoot /inheritance:r /grant:r '*S-1-5-18:(OI)(CI)(F)' '*S-1-5-32-544:(OI)(CI)(F)' '*S-1-5-32-545:(OI)(CI)(RX)' | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw "icacls failed for $InstallRoot"
    }

    & icacls $StateRoot /inheritance:r /grant:r '*S-1-5-18:(OI)(CI)(F)' '*S-1-5-32-544:(OI)(CI)(F)' '*S-1-5-32-545:(OI)(CI)(RX)' | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw "icacls failed for $StateRoot"
    }

    & icacls $LogsRoot /inheritance:r /grant:r '*S-1-5-18:(OI)(CI)(F)' '*S-1-5-32-544:(OI)(CI)(F)' '*S-1-5-32-545:(OI)(CI)(M)' | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw "icacls failed for $LogsRoot"
    }
}

function Start-ActivityWatchTasks {
    param(
        [Parameter(Mandatory = $true)]
        [pscustomobject[]]$TaskDefinitions,
        [string]$RecoveryTaskName = 'ActivityWatch Recovery'
    )

    foreach ($definition in $TaskDefinitions) {
        Start-ScheduledTask -TaskName $definition.LaunchTaskName -ErrorAction SilentlyContinue
    }

    Start-ScheduledTask -TaskName $RecoveryTaskName -ErrorAction SilentlyContinue
}

Export-ModuleMember -Function *-ActivityWatch*, Assert-Administrator, Normalize-ActivityWatchUsers, Get-ActivityWatchPackageUrl, Remove-LegacyActivityWatchEntries