#!/usr/bin/env bash
# sync_data_to_apuana.sh — sync ViTViz datasets to Apuana cluster.
#
# Reads a sweep config YAML, determines which datasets it needs (images_dir +
# Guillaumin masks if FER/mIoU/mAP enabled), checks if they're already on the
# remote, and uploads only what's missing or out-of-sync.
#
# Idempotent: safe to re-run; rsync only transfers changed/new files.
#
# Usage:
#   bash scripts/sync_data_to_apuana.sh <config> <remote_user>@<remote_host> [remote_project_dir]
#
# Example:
#   bash scripts/sync_data_to_apuana.sh \
#       configs/experiments/exp_sweep_a_guillaumin.yaml \
#       lucasddmc@apuana.cin.ufpe.br \
#       ~/ViTViz

set -euo pipefail

if [ "$#" -lt 2 ]; then
  echo "Usage: $0 <config_path> <user@host> [remote_project_dir]"
  echo ""
  echo "  config_path: path to sweep YAML (e.g. configs/experiments/exp_sweep_a_guillaumin.yaml)"
  echo "  user@host:   SSH target (e.g. lucasddmc@apuana.cin.ufpe.br)"
  echo "  remote_project_dir: default ~/ViTViz"
  exit 1
fi

CONFIG_PATH="$1"
SSH_TARGET="$2"
REMOTE_PROJECT_DIR="${3:-~/ViTViz}"

PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
cd "$PROJECT_ROOT"

if [ ! -f "$CONFIG_PATH" ]; then
  echo "ERROR: config not found: $CONFIG_PATH"
  exit 1
fi

# Use Python to inspect the config (more robust than parsing YAML in bash)
INSPECT=$(python3 -c "
import sys, yaml
from pathlib import Path

cfg_path = Path('$CONFIG_PATH')
default = Path('configs/default.yaml')

cfg = {}
if default.exists():
    with open(default) as f: cfg = yaml.safe_load(f) or {}
with open(cfg_path) as f: ovr = yaml.safe_load(f) or {}
for k, v in ovr.items():
    if isinstance(v, dict) and isinstance(cfg.get(k), dict):
        cfg[k].update(v)
    else:
        cfg[k] = v

images_dir = (cfg.get('data') or {}).get('images_dir') or 'data/sample_images'
metrics = (cfg.get('evaluation') or {}).get('selected_metrics') or []
needs_masks = any(m in metrics for m in ('fer', 'miou', 'map'))

print(f'IMAGES_DIR={images_dir}')
print(f'NEEDS_MASKS={1 if needs_masks else 0}')
print(f'METRICS={\",\".join(metrics)}')
")
eval "$INSPECT"

echo "Config: $CONFIG_PATH"
echo "  images_dir:   $IMAGES_DIR"
echo "  needs_masks:  $NEEDS_MASKS (gtsegs_ijcv.mat)"
echo "  metrics:      $METRICS"
echo ""

# Build list of paths to sync
PATHS=("$IMAGES_DIR")
if [ "$NEEDS_MASKS" = "1" ]; then
  PATHS+=("data/gtsegs_ijcv.mat")
fi

# Sync each path
for rel in "${PATHS[@]}"; do
  local_path="$PROJECT_ROOT/$rel"
  remote_path="$REMOTE_PROJECT_DIR/$rel"

  if [ ! -e "$local_path" ]; then
    echo "[$rel]"
    echo "  ERROR: local path missing: $local_path"
    continue
  fi

  # Local stats
  if [ -f "$local_path" ]; then
    local_size=$(stat -f%z "$local_path" 2>/dev/null || stat -c%s "$local_path")
    local_count=1
    echo "[$rel] local file ($((local_size / 1024 / 1024)) MB)"
  else
    local_count=$(find "$local_path" -type f \( -name "*.JPEG" -o -name "*.jpg" -o -name "*.png" -o -name "*.mat" -o -name "*.json" \) 2>/dev/null | wc -l | tr -d ' ')
    echo "[$rel] local dir: $local_count files"
  fi

  # Remote stats (single SSH call) — count + size pra detectar truncamento
  remote_info=$(ssh "$SSH_TARGET" "
    if [ -d '$remote_path' ]; then
      find '$remote_path' -type f \\( -name '*.JPEG' -o -name '*.jpg' -o -name '*.png' -o -name '*.mat' -o -name '*.json' \\) | wc -l
      echo DIR
    elif [ -f '$remote_path' ]; then
      echo 1
      stat -c%s '$remote_path' 2>/dev/null || stat -f%z '$remote_path'
    else
      echo MISSING
      echo MISSING
    fi
  " 2>/dev/null || echo "SSH_ERROR")

  if [ "$remote_info" = "SSH_ERROR" ]; then
    echo "  ERROR: SSH check failed"
    continue
  fi

  remote_count=$(echo "$remote_info" | sed -n '1p' | tr -d ' \n')
  remote_size=$(echo "$remote_info" | sed -n '2p' | tr -d ' \n')

  if [ "$remote_count" = "MISSING" ]; then
    echo "  remote MISSING — uploading..."
  elif [ -f "$local_path" ] && [ "$remote_count" = "1" ] && [ "$remote_size" != "DIR" ] && [ "$remote_size" != "MISSING" ]; then
    # File: comparar bytes pra detectar truncamento
    if [ "$remote_size" = "$local_size" ]; then
      echo "  remote OK ($((remote_size / 1024 / 1024)) MB), skipping"
      continue
    else
      echo "  remote SIZE MISMATCH (local=$((local_size / 1024 / 1024)) MB, remote=$((remote_size / 1024 / 1024)) MB) — re-uploading..."
    fi
  elif [ "$remote_count" -ge "$local_count" ]; then
    echo "  remote OK ($remote_count files), skipping"
    continue
  else
    echo "  remote partial ($remote_count/$local_count) — syncing missing files..."
  fi

  # Ensure parent dir exists
  parent_dir=$(dirname "$remote_path")
  ssh "$SSH_TARGET" "mkdir -p '$parent_dir'"

  # Rsync (resume-friendly)
  if [ -d "$local_path" ]; then
    rsync -azP --info=progress2 "$local_path/" "$SSH_TARGET:$remote_path/"
  else
    rsync -azP --info=progress2 "$local_path" "$SSH_TARGET:$parent_dir/"
  fi

  # Verify
  final_count=$(ssh "$SSH_TARGET" "
    if [ -d '$remote_path' ]; then
      find '$remote_path' -type f \\( -name '*.JPEG' -o -name '*.jpg' -o -name '*.png' -o -name '*.mat' -o -name '*.json' \\) | wc -l
    elif [ -f '$remote_path' ]; then
      echo 1
    else
      echo MISSING
    fi
  " | tr -d ' \n')
  echo "  synced. remote now has $final_count files"
  echo ""
done

echo "Done."