#!/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.
#
# Dos direcciones para la misma máquina: la de LAN se prueba primero porque es
# el caso habitual y es mucho más rápida para rsync; si no responde, se cae a la
# pública. Forzar una con ENLACE_REMOTE_FORCE=lan|wan.
#
#   scripts/remote.sh check                     # conectividad, GPU y entorno
#   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
#   scripts/remote.sh ssh [comando]             # shell o comando suelto
set -euo pipefail
cd "$(dirname "$0")/.."

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

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

# Timeout corto: si no estamos en la LAN, no tiene sentido esperar el TCP.
LAN_TIMEOUT="${ENLACE_LAN_TIMEOUT:-3}"
SSH_OPTS=(-o BatchMode=yes -o StrictHostKeyChecking=accept-new)

_alcanzable() {
    ssh "${SSH_OPTS[@]}" -o "ConnectTimeout=$1" "$2" true 2>/dev/null
}

elegir_host() {
    local lan="${ENLACE_REMOTE_LAN_HOST:-}" wan="${ENLACE_REMOTE_HOST:-}"

    case "${ENLACE_REMOTE_FORCE:-}" in
    lan) echo "${lan:?ENLACE_REMOTE_FORCE=lan pero ENLACE_REMOTE_LAN_HOST no está definido}"; return ;;
    wan) echo "${wan:?ENLACE_REMOTE_FORCE=wan pero ENLACE_REMOTE_HOST no está definido}"; return ;;
    esac

    if [[ -n "$lan" ]] && _alcanzable "$LAN_TIMEOUT" "$lan"; then
        echo "$lan"
        return
    fi
    if [[ -n "$wan" ]]; then
        echo "$wan"
        return
    fi
    echo "definí ENLACE_REMOTE_LAN_HOST o ENLACE_REMOTE_HOST en .env" >&2
    exit 1
}

HOST="$(elegir_host)"
# A stderr para no contaminar la salida de los comandos que se parsean.
if [[ "$HOST" == "${ENLACE_REMOTE_LAN_HOST:-}" ]]; then
    echo "[enlace] servidor: $HOST (LAN)" >&2
else
    echo "[enlace] servidor: $HOST (remoto)" >&2
fi

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

case "$comando" in
check)
    # Todo lo que tiene que estar bien antes de comprometer horas de cómputo o
    # gigabytes de corpus. Falla ruidoso y temprano.
    ssh "${SSH_OPTS[@]}" "$HOST" bash -s -- "$DIR" "$PY" <<'REMOTO'
set -u
DIR="$1"; PY="$2"
echo "host:    $(hostname)"
echo "repo:    $DIR $( [ -d "$DIR/.git" ] && echo '(ok)' || echo '(FALTA: cloná el repo)')"
echo "python:  $( [ -x "$PY" ] && "$PY" --version 2>&1 || echo 'FALTA el venv')"
echo "tmux:    $(command -v tmux >/dev/null && tmux -V || echo 'FALTA (las corridas largas lo necesitan)')"
if command -v nvidia-smi >/dev/null; then
    nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader | sed 's/^/gpu:     /'
else
    echo "gpu:     nvidia-smi no está en PATH"
fi
[ -x "$PY" ] && "$PY" - <<'PY'
try:
    import torch
except ImportError:
    print("torch:   no instalado (corré `remote.sh sync`)"); raise SystemExit
print(f"torch:   {torch.__version__} | cuda disponible: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    cap = torch.cuda.get_device_capability(0)
    print(f"         sm_{cap[0]}{cap[1]} | bf16: {torch.cuda.is_bf16_supported()}")
    print(f"         perfil sugerido: {'blackwell-5090' if cap >= (12,0) else 'turing-2060'}")
PY
REMOTO
    ;;

sync)
    git push
    ssh "${SSH_OPTS[@]}" "$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 "${SSH_OPTS[@]}" -t "$HOST" "cd '$DIR' && mkdir -p runs && \
        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 "${SSH_OPTS[@]}" "$HOST" "cd '$DIR' && ls -t runs/*.log | head -1 | xargs tail -f"
    else
        ssh "${SSH_OPTS[@]}" "$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 "${SSH_OPTS[@]}" "$HOST" "nvidia-smi"
    ;;

ssh)
    if [[ $# -eq 0 ]]; then
        ssh "${SSH_OPTS[@]}" -t "$HOST" "cd '$DIR' && exec \$SHELL -l"
    else
        ssh "${SSH_OPTS[@]}" "$HOST" "cd '$DIR' && $*"
    fi
    ;;

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