#!/usr/bin/env bash
# Trabajo contra el servidor de entrenamiento por SSH.
#
# El código viaja por git; los datos, las bases y los checkpoints viven en el
# servidor y no se descargan enteros. Solo vuelven artefactos chicos: métricas,
# muestras generadas y logs.
#
#   scripts/remote.sh sync                      # empuja el código al servidor
#   scripts/remote.sh run <config.yaml> [args]  # entrena bajo tmux
#   scripts/remote.sh logs [corrida]            # sigue el log en vivo
#   scripts/remote.sh pull <corrida>            # trae métricas y muestras
#   scripts/remote.sh gpu                       # estado de la GPU
set -euo pipefail
cd "$(dirname "$0")/.."

[[ -f .env ]] && set -a && source .env && set +a

HOST="${ENLACE_REMOTE_HOST:?definí ENLACE_REMOTE_HOST en .env (ver .env.example)}"
DIR="${ENLACE_REMOTE_DIR:?definí ENLACE_REMOTE_DIR en .env}"
PY="${ENLACE_REMOTE_PYTHON:-$DIR/.venv/bin/python}"

comando="${1:-}"
shift || true

case "$comando" in
sync)
    git push
    ssh "$HOST" "cd '$DIR' && git pull --ff-only && $PY -m pip install -q -e ."
    ;;

run)
    config="${1:?uso: remote.sh run <config.yaml> [overrides...]}"
    shift
    # El nombre de la sesión sale del nombre de la config, así que dos corridas
    # distintas no se pisan y `logs` sabe a cuál conectarse.
    sesion="enlace-$(basename "$config" .yaml)"
    # tmux es lo que hace que cortar el SSH no mate el entrenamiento.
    ssh -t "$HOST" "cd '$DIR' && tmux new-session -d -s '$sesion' \
        \"$PY -m enlace.train.train '$config' $* 2>&1 | tee -a 'runs/$sesion.log'\" \
        && echo 'corriendo en tmux: $sesion'"
    ;;

logs)
    sesion="${1:-}"
    if [[ -z "$sesion" ]]; then
        ssh "$HOST" "cd '$DIR' && ls -t runs/*.log | head -1 | xargs tail -f"
    else
        ssh "$HOST" "cd '$DIR' && tail -f 'runs/$sesion.log'"
    fi
    ;;

pull)
    corrida="${1:?uso: remote.sh pull <corrida>}"
    mkdir -p "runs/$corrida"
    # Solo métricas y muestras: los checkpoints se quedan en el servidor.
    rsync -av --include='metrics.jsonl' --include='samples.txt' --include='*.yaml' \
        --exclude='*' "$HOST:$DIR/runs/$corrida/" "runs/$corrida/"
    ;;

gpu)
    ssh "$HOST" "nvidia-smi"
    ;;

*)
    sed -n '2,14p' "$0"
    exit 2
    ;;
esac