# Cognio 远程构建 Docker 镜像脚本

[Console]::OutputEncoding = [System.Text.Encoding]::UTF8

# ============================================================
# 配置区
# ============================================================

# 远程服务器配置
$remoteUser = "root"
$remoteHost = "192.168.1.99"
$remotePort = "22"

# Git 仓库配置
$repoUrl = "http://ds:39418/AI/Cognio"
$repoDir = "Cognio"

# Docker 镜像配置
$imageName = "amwpfiqvy/cognio"

# Docker 镜像加速源（随机选择一个，不带 https:// 前缀）
$dockerMirrors = @(
    "docker.m.daocloud.io"
    "docker.1panel.live"
    "hub.rat.dev"
)
$dockerMirror = $dockerMirrors | Get-Random

# 构建时代理配置（用于 poetry/pip 下载依赖）
$buildProxy = "http://192.168.1.222:7890"

# tmux 会话名
$tmuxSession = "dbx"

# ============================================================
# 主程序
# ============================================================

Write-Host ""
Write-Host "========================================" -ForegroundColor Cyan
Write-Host "  Cognio 远程构建 Docker 镜像" -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan
Write-Host ""

# Docker 镜像标签：提示用户输入（默认包含 latest，-l 排除 latest）
$imageTagInput = Read-Host "请输入额外的 Docker 镜像标签 (多个用空格分隔，-l=排除latest，留空则仅推送 latest)"

$inputTags = @()
$excludeLatest = $false

if ($imageTagInput -ne "") {
    $inputTags = $imageTagInput -split '\s+' | Where-Object { $_ -ne "" }
    if ($inputTags -contains "-l") {
        $excludeLatest = $true
        $inputTags = $inputTags | Where-Object { $_ -ne "-l" }
    }
}

# 构建最终标签列表
if ($excludeLatest) {
    $imageTags = $inputTags
} else {
    $imageTags = @("latest") + $inputTags | Select-Object -Unique
}

# 构建 -t 参数列表
$tagParams = ($imageTags | ForEach-Object { "-t ${imageName}:$_" }) -join " "

Write-Host "镜像标签: $($imageTags -join ', ')" -ForegroundColor Yellow
Write-Host "镜像加速: $dockerMirror" -ForegroundColor Yellow
Write-Host ""

# 支持命令行参数覆盖默认配置
if ($args.Count -ge 1) { $remoteUser = $args[0] }
if ($args.Count -ge 2) { $remoteHost = $args[1] }
if ($args.Count -ge 3) { $remotePort = $args[2] }

Write-Host "连接到: $remoteUser@$remoteHost`:$remotePort" -ForegroundColor Yellow
Write-Host ""

# ============================================================
# 构建远程执行的 Shell 脚本
# ============================================================

# 生成远程脚本内容（写入临时文件执行，避免引号嵌套问题）
$scriptContent = @"
#!/bin/bash
cd ~
if [ -d "$repoDir" ]; then
    echo "目录已存在，还原本地修改并拉取最新代码..."
    cd $repoDir
    git checkout .
    git pull
else
    echo "目录不存在，正在克隆仓库..."
    git clone $repoUrl
    cd $repoDir
fi
echo "修改 Dockerfile 使用镜像加速..."
sed -i '1s|^FROM python|FROM \$\{DOCKER_MIRROR\}/library/python|' Dockerfile
sed -i '1i ARG DOCKER_MIRROR=docker.io' Dockerfile
echo "添加 apt 国内镜像源..."
sed -i '/^FROM/a RUN sed -i "s|http://deb.debian.org|https://mirrors.aliyun.com|g" /etc/apt/sources.list.d/debian.sources 2>/dev/null || sed -i "s|http://deb.debian.org|https://mirrors.aliyun.com|g" /etc/apt/sources.list 2>/dev/null || true' Dockerfile
echo "配置 pip 国内镜像源..."
sed -i 's|pip install|pip install -i https://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com|g' Dockerfile
echo "配置 Poetry 国内镜像源（环境变量）..."
sed -i '/^FROM/a ENV PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/ PIP_TRUSTED_HOST=mirrors.aliyun.com' Dockerfile
echo "配置 pip 只使用预编译包（避免 triton 编译问题）..."
sed -i '/^FROM/a ENV PIP_ONLY_BINARY=:all:' Dockerfile
echo "设置代理环境变量（用于推送镜像）..."
export HTTP_PROXY=$buildProxy
export HTTPS_PROXY=$buildProxy
echo "配置 Docker daemon 代理（用于 buildx 推送）..."
sudo mkdir -p /etc/systemd/system/docker.service.d
cat << PROXYEOF | sudo tee /etc/systemd/system/docker.service.d/proxy.conf > /dev/null
[Service]
Environment="HTTP_PROXY=$buildProxy"
Environment="HTTPS_PROXY=$buildProxy"
Environment="NO_PROXY=localhost,127.0.0.1,mirrors.aliyun.com"
PROXYEOF
sudo systemctl daemon-reload
sudo systemctl restart docker
sleep 3
echo "临时禁用 IPv6（避免推送超时）..."
sudo sysctl -w net.ipv6.conf.all.disable_ipv6=1
sudo sysctl -w net.ipv6.conf.default.disable_ipv6=1
echo "创建带代理的 buildx builder..."
docker buildx rm mybuilder 2>/dev/null || true
docker buildx create --name mybuilder --driver docker-container \
    --driver-opt env.HTTP_PROXY=$buildProxy \
    --driver-opt env.HTTPS_PROXY=$buildProxy \
    --use
docker buildx inspect mybuilder --bootstrap
echo "验证 builder 状态..."
docker buildx ls
echo "正在构建并推送 Docker 镜像..."
echo "使用镜像加速: $dockerMirror"
echo "使用代理: $buildProxy"
docker buildx build --builder mybuilder \
    --platform linux/amd64,linux/arm64 \
    --build-arg DOCKER_MIRROR=$dockerMirror \
    --build-arg HTTP_PROXY=$buildProxy \
    --build-arg HTTPS_PROXY=$buildProxy \
    $tagParams \
    --cache-from type=local,src=`$HOME/.buildx-cache \
    --cache-to type=local,dest=`$HOME/.buildx-cache,mode=max \
    --push \
    .
echo "恢复 IPv6..."
sudo sysctl -w net.ipv6.conf.all.disable_ipv6=0
sudo sysctl -w net.ipv6.conf.default.disable_ipv6=0
"@

# 将脚本内容转为 base64，避免特殊字符问题
$scriptBase64 = [Convert]::ToBase64String([System.Text.Encoding]::UTF8.GetBytes($scriptContent))

# tmux 命令：解码并执行脚本
$remoteCmd = "tmux has-session -t $tmuxSession 2>/dev/null || tmux new-session -d -s $tmuxSession; " +
             "tmux send-keys -t $tmuxSession 'echo $scriptBase64 | base64 -d > /tmp/build_cognio.sh && bash /tmp/build_cognio.sh' Enter"

# ============================================================
# 执行远程命令
# ============================================================

Write-Host "正在发送命令到 tmux 会话: $tmuxSession" -ForegroundColor Yellow
Write-Host ""

& ssh -p $remotePort "$remoteUser@$remoteHost" $remoteCmd

if ($LASTEXITCODE -eq 0) {
    Write-Host ""
    Write-Host "========================================" -ForegroundColor Green
    Write-Host "  命令已发送到 tmux 会话: $tmuxSession" -ForegroundColor Green
    Write-Host "========================================" -ForegroundColor Green
    Write-Host ""
    Write-Host "推送目标:" -ForegroundColor Yellow
    foreach ($tag in $imageTags) {
        Write-Host "  - ${imageName}:${tag}" -ForegroundColor Cyan
    }
    Write-Host ""
    Write-Host "查看构建进度:" -ForegroundColor Yellow
    Write-Host "  ssh -t -p $remotePort $remoteUser@$remoteHost `"tmux attach -t $tmuxSession`"" -ForegroundColor Cyan
} else {
    Write-Host ""
    Write-Host "部署失败，请检查错误信息" -ForegroundColor Red
}

Write-Host ""
pause