# Copyright (c) 2026 Progralink. Licensed under the DevCage License — see LICENSE.
#
# Container.ps1 — Container build and run logic for the DevCage PowerShell
# launcher. Dot-sourced by devcage.ps1.
# Depends on: $ContainerEngine, $ScriptDir, $ComposeFile, $Cwd, $WorkspaceSlug,
#             $ContainerName, $EmbedDocker, $GpuType, $NoGpu, $NoReuse,
#             $Recreate, $IdleTimeout, $InteractiveCommands,
#             $script:RunArgs, $script:ResolvedMemory, $script:ResolvedCpus,
#             $script:DirCacheCount, $script:AllVolumeArgs, $script:ExecEnvArgs,
#             $DnsServers, $BindAddr, $PortList.

# ── Compose file stacking and progress flags ──────────────────────
# Sets $script:ComposeFiles, $script:BuildProgressArgs, $script:RunProgressArgs.
function Initialize-ComposeFiles {
    $script:ComposeFiles = @("-f", $ComposeFile)
    if (-not $NoGpu -and $GpuType) {
        switch ($GpuType) {
            "nvidia" {
                $gpuOverlay = Join-Path (Join-Path $ScriptDir "sandbox") "docker-compose.gpu-nvidia.yml"
                if (Test-Path $gpuOverlay) {
                    $script:ComposeFiles += @("-f", $gpuOverlay)
                    Write-Info "GPU: NVIDIA runtime detected - enabling GPU passthrough."
                }
            }
            "amd" {
                $gpuOverlay = Join-Path (Join-Path $ScriptDir "sandbox") "docker-compose.gpu-amd.yml"
                if (Test-Path $gpuOverlay) {
                    $script:ComposeFiles += @("-f", $gpuOverlay)
                    Write-Info "GPU: AMD runtime detected - enabling GPU passthrough."
                }
            }
        }
    }
    # Docker-only: Podman compose variants may not support --progress.
    $script:BuildProgressArgs = @()
    $script:RunProgressArgs = @()
    if ($ContainerEngine -eq "docker") {
        $script:BuildProgressArgs = @("--progress", "quiet")
        $script:RunProgressArgs = @("--progress", "quiet")
    }
}

# ── Build basic run args ──────────────────────────────────────────
# Initializes $script:RunArgs with core flags (reuse labels, apt/brew/snap,
# verbose, yolo-config, project-init, project-crontab signals).
# Note: session log, env vars, ports, and env files are added separately.
function Build-RunArgs {
    $script:RunArgs = @()
    if ($NoReuse) {
        $script:RunArgs += @("--rm")
    }
    # Tag auto-mode resources for fair-share rebalancing across concurrent instances
    if ($MemoryLimit -eq 'auto') {
        $script:RunArgs += @("--label", "dev-sandbox.memory-mode=auto")
    }
    if ($CpuLimit -eq 'auto') {
        $script:RunArgs += @("--label", "dev-sandbox.cpu-mode=auto")
    }
    if ($AptPackages) {
        $script:RunArgs += @("-e", "APT_PACKAGES=$AptPackages")
    }
    if ($BrewPackages) {
        $script:RunArgs += @("-e", "BREW_PACKAGES=$BrewPackages")
    }
    if ($SnapPackages) {
        $script:RunArgs += @("-e", "SNAP_PACKAGES=$SnapPackages")
    }
    if ($Verbose) {
        $script:RunArgs += @("-e", "SANDBOX_VERBOSE=1")
    }
    if ($YoloConfig) {
        $script:RunArgs += @("-e", "SANDBOX_YOLO_CONFIG=1")
    }
    if ($SandboxProjectInit) {
        $script:RunArgs += @("-e", "SANDBOX_PROJECT_INIT=1")
    }
    if ($SandboxProjectCrontab) {
        $script:RunArgs += @("-e", "SANDBOX_PROJECT_CRONTAB=1")
    }
}

# ── Port forwarding args ──────────────────────────────────────────
# Validates and appends --publish args to $script:RunArgs.
function Build-PortArgs {
    if (-not $PortList) { return }
    $seenPorts = @{}
    foreach ($portSpec in $PortList -split ',') {
        $portSpec = $portSpec.Trim()
        $colonCount = ($portSpec -split ':').Count - 1

        if ($colonCount -eq 0) {
            # Simple port number: 3000 → BindAddr:3000:3000
            $portNum = 0
            if ($portSpec -match '^\d+$' -and $portSpec.Length -le 5 -and [int]::TryParse($portSpec, [ref]$portNum) -and $portNum -ge 1 -and $portNum -le 65535) {
                $portKey = "${BindAddr}:${portSpec}:${portSpec}"
                if (-not $seenPorts.ContainsKey($portKey)) {
                    $seenPorts[$portKey] = $true
                    $script:RunArgs += @("--publish", $portKey)
                } else {
                    Write-Host "Warning: duplicate port mapping $portKey ignored" -ForegroundColor Yellow
                }
            } else {
                Write-Host "Error: invalid port '$portSpec' (must be 1-65535)" -ForegroundColor Red
                Remove-SandboxTmpDir; exit 1
            }
        } elseif ($colonCount -eq 1) {
            # Host:container mapping: 5080:80 → BindAddr:5080:80
            $parts = $portSpec -split ':'
            $hostPort = 0; $containerPort = 0
            if ($parts[0] -match '^\d+$' -and [int]::TryParse($parts[0], [ref]$hostPort) -and $hostPort -ge 1 -and $hostPort -le 65535 -and
                $parts[1] -match '^\d+$' -and [int]::TryParse($parts[1], [ref]$containerPort) -and $containerPort -ge 1 -and $containerPort -le 65535) {
                $portKey = "${BindAddr}:${hostPort}:${containerPort}"
                if (-not $seenPorts.ContainsKey($portKey)) {
                    $seenPorts[$portKey] = $true
                    $script:RunArgs += @("--publish", $portKey)
                } else {
                    Write-Host "Warning: duplicate port mapping $portKey ignored" -ForegroundColor Yellow
                }
            } else {
                Write-Host "Error: invalid port mapping '$portSpec' (format: HOST_PORT:CONTAINER_PORT, both 1-65535)" -ForegroundColor Red
                Remove-SandboxTmpDir; exit 1
            }
        } elseif ($colonCount -eq 2) {
            # Full Docker syntax: 0.0.0.0:5080:80 → use as-is
            $parts = $portSpec -split ':'
            $hostPort = 0; $containerPort = 0
            if ((Test-IPv4Address $parts[0]) -and
                $parts[1] -match '^\d+$' -and [int]::TryParse($parts[1], [ref]$hostPort) -and $hostPort -ge 1 -and $hostPort -le 65535 -and
                $parts[2] -match '^\d+$' -and [int]::TryParse($parts[2], [ref]$containerPort) -and $containerPort -ge 1 -and $containerPort -le 65535) {
                $portKey = $portSpec
                if (-not $seenPorts.ContainsKey($portKey)) {
                    $seenPorts[$portKey] = $true
                    $script:RunArgs += @("--publish", $portKey)
                } else {
                    Write-Host "Warning: duplicate port mapping $portKey ignored" -ForegroundColor Yellow
                }
            } else {
                Write-Host "Error: invalid port mapping '$portSpec' (format: IP:HOST_PORT:CONTAINER_PORT)" -ForegroundColor Red
                Remove-SandboxTmpDir; exit 1
            }
        } else {
            Write-Host "Error: invalid port mapping '$portSpec' (too many colons)" -ForegroundColor Red
            Remove-SandboxTmpDir; exit 1
        }
    }
}

# ── Detect interactive vs non-interactive command ─────────────────
# Sets $script:InteractiveCmd.
function Detect-InteractiveCmd {
    $script:InteractiveCmd = $true
    if ($ContainerArgs.Count -gt 0) {
        $firstCmd = Split-Path -Leaf $ContainerArgs[0]
        if ($firstCmd -notin $InteractiveCommands) {
            $script:InteractiveCmd = $false
        }
    }
}

# ── Common docker run args ────────────────────────────────────────
# Sets $script:DrArgs (flags shared by single-shot and reuse-new-container paths).
function Build-CommonRunArgs {
    $dnsList = if ($DnsServers) { $DnsServers } else { "8.8.8.8,1.1.1.1" }
    $script:DrArgs = @(
        "--hostname", $WorkspaceSlug,
        "--label", "dev-sandbox.workspace=$($env:WORKSPACE)",
        "--label", "dev-sandbox.version=$DevcageVersion",
        "--cap-drop", "ALL",
        "--cap-add", "SETUID", "--cap-add", "SETGID", "--cap-add", "SETPCAP",
        "--cap-add", "CHOWN", "--cap-add", "DAC_OVERRIDE", "--cap-add", "FOWNER",
        "--cap-add", "NET_ADMIN", "--cap-add", "NET_RAW",
        "--security-opt", "no-new-privileges:true",
        "--memory", $env:SANDBOX_MEMORY,
        "--memory-swap", $env:SANDBOX_MEMORY_SWAP,
        "--cpus", $env:SANDBOX_CPUS,
        "--pids-limit", "4096",
        "--pull", "never",
        "-e", "TERM=$(if ($env:TERM) { $env:TERM } else { 'xterm-256color' })",
        "-e", "COLORTERM=$(if ($env:COLORTERM) { $env:COLORTERM } else { '' })",
        "-e", "HOST_UID=$($env:HOST_UID)",
        "-e", "HOST_WORKSPACE=$($env:WORKSPACE)",
        "-e", "AGENT_CONFIG_RO=$($env:AGENT_CONFIG_RO)",
        "-e", "SANDBOX_VERBOSE=$(if ($env:SANDBOX_VERBOSE) { $env:SANDBOX_VERBOSE } else { '0' })",
        "-e", "DISABLE_TELEMETRY=1",
        "-e", "DO_NOT_TRACK=1",
        "-v", "$($env:WORKSPACE):/workspace"
    )
    # DNS servers
    foreach ($dip in ($dnsList -split ',')) {
        $dip = $dip.Trim()
        if ($dip) { $script:DrArgs += @("--dns", $dip) }
    }
}

# ── Build exec env args (API keys + per-session vars) ──────────
# Sets $script:ExecEnvArgs for docker exec calls.
function Build-ExecEnvArgs {
    $script:ExecEnvArgs = @()
    foreach ($apiVar in $ApiKeyVarsExec) {
        $val = [Environment]::GetEnvironmentVariable($apiVar)
        if ($val) { $script:ExecEnvArgs += @("-e", $apiVar) }
    }
    $script:ExecEnvArgs += @("-e", "TERM=$(if ($env:TERM) { $env:TERM } else { 'xterm-256color' })")
    $script:ExecEnvArgs += @("-e", "COLORTERM=$(if ($env:COLORTERM) { $env:COLORTERM } else { 'truecolor' })")
    if ($SessionLog) {
        $script:ExecEnvArgs += @("-e", "DEVCAGE_SESSION_LOG=/workspace/.devcage/logs/session-${SessionTs}.log")
    }
}

# ── Assemble all volume args ──────────────────────────────────────
# Sets $script:AllVolumeArgs (for docker run) and $script:AllComposeVolumeArgs
# (for docker compose run — excludes --mount which compose doesn't support).
function Build-AllVolumeArgs {
    $script:AllVolumeArgs = @()
    $script:AllVolumeArgs += $VolumeArgs
    $script:AllVolumeArgs += $InitVolumeArgs
    $script:AllVolumeArgs += $CrontabVolumeArgs
    $script:AllVolumeArgs += $AgentsVolumeArgs
    $script:AllVolumeArgs += $VersionsVolumeArgs
    $script:AllVolumeArgs += $CacheVolumeArgs
    $script:AllVolumeArgs += $DirCacheVolumeArgs
    $script:AllVolumeArgs += $CustomVolumeArgs

    # Compose-compatible: excludes --mount (not supported by docker compose run).
    $script:AllComposeVolumeArgs = @()
    $script:AllComposeVolumeArgs += $VolumeArgs
    $script:AllComposeVolumeArgs += $InitVolumeArgs
    $script:AllComposeVolumeArgs += $CrontabVolumeArgs
    $script:AllComposeVolumeArgs += $AgentsVolumeArgs
    $script:AllComposeVolumeArgs += $VersionsVolumeArgs
    $script:AllComposeVolumeArgs += $CacheVolumeArgs
    # Filter --mount args from DirCacheVolumeArgs
    $skipNext = $false
    foreach ($arg in $DirCacheVolumeArgs) {
        if ($skipNext) { $skipNext = $false; continue }
        if ($arg -eq "--mount") { $skipNext = $true; continue }
        $script:AllComposeVolumeArgs += $arg
    }
    $script:AllComposeVolumeArgs += $CustomVolumeArgs
}

# ── Container reuse logic ─────────────────────────────────────────
# Discovers existing container, attaches via exec, or sets up for new creation.
# Sets $script:ReuseAttached = $true if successfully attached to existing container.
function Handle-ContainerReuse {
    $script:ReuseAttached = $false

    if ($NoReuse -or $EmbedDocker) { return }

    # Handle --recreate
    if ($Recreate) {
        $oldCid = & $ContainerEngine ps -aq --filter "label=dev-sandbox.workspace=$($env:WORKSPACE)" 2>$null | Select-Object -First 1
        if ($oldCid) {
            Write-Info "Recreating container..."
            & $ContainerEngine rm -f $oldCid *>$null
        }
    }

    # Discover existing container
    $existing = & $ContainerEngine ps -q --filter "label=dev-sandbox.workspace=$($env:WORKSPACE)" 2>$null | Select-Object -First 1
    $stopped = $null
    if (-not $existing) {
        $stopped = & $ContainerEngine ps -aq --filter "status=exited" --filter "label=dev-sandbox.workspace=$($env:WORKSPACE)" 2>$null | Select-Object -First 1
    }

    # Validate running container version
    if ($existing) {
        $containerVersion = $null
        try {
            $containerVersion = (& $ContainerEngine inspect --format '{{json .Config.Labels}}' $existing 2>$null | ConvertFrom-Json).'dev-sandbox.version'
        } catch {
            # Corrupt or partial JSON from docker inspect — treat as version mismatch (safe: triggers recreation)
        }
        if ($containerVersion -ne $DevcageVersion) {
            Write-Info "Container running older DevCage version ($containerVersion). Stopping and recreating..."
            & $ContainerEngine rm -f $existing *>$null
            $existing = $null
        }
    }

    if ($existing) {
        # Attach to running container via docker exec
        if ($PortList) {
            $containerPorts = (& $ContainerEngine inspect --format '{{json .NetworkSettings.Ports}}' $existing 2>$null)
            foreach ($rp in ($PortList -split ',')) {
                $rp = $rp.Trim()
                $cp = ($rp -split ':')[-1]
                if ($containerPorts -notmatch "${cp}/tcp") {
                    Write-Warning "Port $cp not available in running container."
                    Write-Warning "  Use ``devcage --recreate -p $PortList ...`` to rebuild with additional ports."
                }
            }
        }

        $ndc = @(& $ContainerEngine ps -q --filter "label=dev-sandbox.workspace" 2>$null).Count - 1
        $wsSessions = [int](& $ContainerEngine exec $existing sh -c 'count=0; for f in /tmp/.devcage-sessions/[0-9]*; do [ -f "$f" ] || continue; p="${f##*/}"; if [ -d "/proc/$p" ]; then count=$(( count + 1 )); else rm -f "$f" 2>/dev/null; fi; done; echo "$count"' 2>$null)
        if ($ndc -gt 0 -or $wsSessions -gt 0) {
            $msg = @()
            if ($ndc -gt 0) {
                $msg += if ($ndc -eq 1) { "There is 1 other DevCage container running." } else { "There are $ndc other DevCage containers running." }
            }
            if ($wsSessions -gt 0) {
                $msg += if ($wsSessions -eq 1) { "There is 1 other session in this workspace." } else { "There are $wsSessions other sessions in this workspace." }
            }
            Write-Info ($msg -join " ")
        }
        $dindLabel = if ($EmbedDocker) { "on" } else { "off" }
        $statusLine = "Engine: $ContainerEngine | DinD: $dindLabel"
        if ($DirCacheCount -gt 0) { $statusLine += " | CacheDirs: $DirCacheCount" }
        $statusLine += " | Memory: $ResolvedMemory | CPUs: $ResolvedCpus | Container: reused"
        Write-Info $statusLine

        $script:ReuseContainerId = $existing
        $ExecTtyArgs = if ($InteractiveCmd) { @("-i", "-t") } else { @("-i") }
        $allExecArgs = @("exec") + $ExecTtyArgs + $ExecEnvArgs + @($existing, "/usr/local/bin/devcage-session") + $ContainerArgs
        & $ContainerEngine @allExecArgs
        $script:ReuseAttached = $true
        return
    }

    if ($stopped) {
        # Stopped container found — remove and recreate with fresh config
        Write-Info "Recreating stopped container..."
        & $ContainerEngine rm $stopped *>$null
        $stopped = $null
    }

    # No existing container — create new persistent container
    if (-not $script:ReuseAttached -and -not $existing -and -not $stopped) {
        $ndc = @(& $ContainerEngine ps -q --filter "label=dev-sandbox.workspace" 2>$null).Count
        if ($ndc -gt 0) {
            Write-Info $(if ($ndc -eq 1) { "There is 1 other DevCage container running." } else { "There are $ndc other DevCage containers running." })
        }
        $dindLabel = if ($EmbedDocker) { "on" } else { "off" }
        $statusLine = "Engine: $ContainerEngine | DinD: $dindLabel"
        if ($DirCacheCount -gt 0) { $statusLine += " | CacheDirs: $DirCacheCount" }
        $statusLine += " | Memory: $ResolvedMemory | CPUs: $ResolvedCpus | Container: new | Initializing..."
        Write-Info $statusLine

        # API keys NOT passed at docker run in reuse mode — they are passed
        # per-session via docker exec -e, avoiding persistent exposure in
        # docker inspect. The entrypoint does not need API keys.

        # Reuse mode: create detached container
        $script:DrArgs += @(
            "--name", $ContainerName,
            "-d",
            "--restart", "on-failure",
            "-e", "DEVCAGE_REUSE=1",
            "-e", "DEVCAGE_IDLE_TIMEOUT=$IdleTimeout"
        )

        $newCid = & $ContainerEngine run @DrArgs @RunArgs @AllVolumeArgs $ImageName
        if ($LASTEXITCODE -ne 0) {
            Write-Host "Error: failed to create container" -ForegroundColor Red
            exit 1
        }
        $script:ReuseContainerId = $newCid

        # Wait for healthcheck
        $waitCount = 0
        while ($waitCount -lt 240) {
            $prevEAP = $ErrorActionPreference; $ErrorActionPreference = "SilentlyContinue"
            & $ContainerEngine exec $newCid test -f /tmp/.devcage-ready *>$null
            $ErrorActionPreference = $prevEAP
            if ($LASTEXITCODE -eq 0) { break }
            # Check if container died during init
            $state = & $ContainerEngine inspect --format '{{.State.Status}}' $newCid 2>$null
            if ($state -eq "exited" -or $state -eq "dead") {
                Write-Host "Error: container exited during initialization" -ForegroundColor Red
                & $ContainerEngine logs $newCid 2>&1
                & $ContainerEngine rm $newCid *>$null
                exit 1
            }
            Start-Sleep -Milliseconds 500
            $waitCount++
        }
        if ($waitCount -ge 240) {
            Write-Warning "Container did not become ready within 120s -- proceeding anyway"
        }

        # Attach via exec
        $ExecTtyArgs = if ($InteractiveCmd) { @("-i", "-t") } else { @("-i") }
        $allExecArgs = @("exec") + $ExecTtyArgs + $ExecEnvArgs + @($newCid, "/usr/local/bin/devcage-session") + $ContainerArgs
        & $ContainerEngine @allExecArgs
        $script:ReuseAttached = $true
    }
}

# ── Single-shot container mode (--no-reuse or --embed-docker) ─────
function Run-SingleShotContainer {
    $ndc = @(& $ContainerEngine ps -q --filter "label=dev-sandbox.workspace" 2>$null).Count
    if ($ndc -gt 0) {
        Write-Info $(if ($ndc -eq 1) { "There is 1 other DevCage container running." } else { "There are $ndc other DevCage containers running." })
    }
    $dindLabel = if ($EmbedDocker) { "on" } else { "off" }
    $statusLine = "Engine: $ContainerEngine | DinD: $dindLabel"
    if ($DirCacheCount -gt 0) { $statusLine += " | CacheDirs: $DirCacheCount" }
    $statusLine += " | Memory: $ResolvedMemory | CPUs: $ResolvedCpus | Initializing..."
    Write-Info $statusLine

    # Export DNS env vars for compose variable substitution
    $dnsCompose = if ($DnsServers) { $DnsServers } else { "8.8.8.8,1.1.1.1" }
    $dnsComposeArr = $dnsCompose -split ','
    $env:DNS_1 = if ($dnsComposeArr.Count -ge 1) { $dnsComposeArr[0].Trim() } else { "8.8.8.8" }
    $env:DNS_2 = if ($dnsComposeArr.Count -ge 2) { $dnsComposeArr[1].Trim() } else { "1.1.1.1" }

    # DinD LAN isolation
    $DindLanBlock = Join-Path (Join-Path $ScriptDir "sandbox") "dind-lan-block.sh"
    if ($EmbedDocker -and (Test-Path $DindLanBlock)) {
        $prevEAP = $ErrorActionPreference; $ErrorActionPreference = "SilentlyContinue"
        & $ContainerEngine compose @RunProgressArgs @ComposeFiles up -d --wait dind *>$null
        $dindUpOk = $LASTEXITCODE -eq 0
        if ($dindUpOk) {
            $ErrorActionPreference = "Continue"
            Get-Content $DindLanBlock -Raw | & $ContainerEngine compose @ComposeFiles exec -T dind sh -c "tr -d '\015' | sh -s" 2>&1
            if ($LASTEXITCODE -ne 0) {
                Write-Warning "DinD LAN isolation failed - containers inside DinD may reach LAN addresses"
                Write-Warning "  Tip: omit --embed-docker if Docker-in-Docker is not needed"
            }
        } else {
            Write-Warning "DinD sidecar failed to start - containers inside DinD may reach LAN addresses"
            Write-Warning "  Tip: omit --embed-docker if Docker-in-Docker is not needed"
        }
        $ErrorActionPreference = $prevEAP
    }

    if (-not $EmbedDocker) {
        # Fast path: bypass docker compose for single-container mode
        foreach ($apiVar in $ApiKeyVarsCore) {
            $val = [Environment]::GetEnvironmentVariable($apiVar)
            if ($val) { $script:DrArgs += @("-e", $apiVar) }
        }
        if ($InteractiveCmd) {
            $script:DrArgs += @("-i", "-t")
        } else {
            $script:DrArgs += @("-i")
        }
        if ($NoReuse) {
            $script:DrArgs += @("-e", "DEVCAGE_REUSE=0")
        }
        & $ContainerEngine run @DrArgs @RunArgs @AllVolumeArgs $ImageName @ContainerArgs
    } else {
        # DinD enabled: use docker compose for two-container orchestration
        $ComposeExtra = @("--quiet-pull")
        if (-not $InteractiveCmd) {
            $ComposeExtra += @("-T")
        }
        # Compose always uses --rm (reuse mode not supported with DinD)
        & $ContainerEngine compose @RunProgressArgs @ComposeFiles run --rm @ComposeExtra @RunArgs @AllComposeVolumeArgs workspace @ContainerArgs
    }
}

# ── Run container (orchestrates reuse and single-shot) ────────────
# Sets $script:RunExitCode.
function Run-Container {
    $script:ReuseContainerId = $null  # set by reuse path; used for OOM-vs-idle diagnostics
    Handle-ContainerReuse
    if (-not $ReuseAttached) {
        Run-SingleShotContainer
    }
    $script:RunExitCode = $LASTEXITCODE
}
