#!/bin/bash
# OpenTrack runtime initialization script
# This script downloads datasets and sets up symlinks at container startup

set -e  # Exit on error

echo "🤖 Initializing OpenTrack workspace..."
echo "======================================"

# Define paths
WORKSPACE="/data/workspaces/opentrack"
DATASETS_DIR="$WORKSPACE/datasets"
MODELS_DIR="$WORKSPACE/models"
VIDEOS_DIR="$WORKSPACE/videos"
OPENTRACK_REPO="$HOME/OpenTrack"

# Create workspace directories if they don't exist
mkdir -p "$DATASETS_DIR"
mkdir -p "$MODELS_DIR"
mkdir -p "$VIDEOS_DIR"
echo "✓ Workspace directories ready"

# Download mocap datasets if not already present
MOCAP_DIR="$DATASETS_DIR/lafan1/UnitreeG1"
mkdir -p "$MOCAP_DIR"

if [ -z "$(ls -A "$MOCAP_DIR" 2>/dev/null)" ]; then
    echo ""
    echo "📥 Downloading motion capture datasets..."
    echo "   This may take a few minutes on first run..."

    # Use Python with huggingface_hub to download
    python3 << 'EOF'
import sys
from pathlib import Path

try:
    from huggingface_hub import snapshot_download

    datasets_dir = Path("$DATASETS_DIR")

    print("   Downloading from robfiras/loco-mujoco-datasets...")
    snapshot_path = snapshot_download(
        repo_id="robfiras/loco-mujoco-datasets",
        repo_type="dataset",
        allow_patterns="Lafan1/mocap/UnitreeG1/*.npz",
        local_dir=str(datasets_dir),
        local_dir_use_symlinks=False
    )

    # Count downloaded files
    mocap_dir = datasets_dir / "lafan1" / "UnitreeG1"
    npz_files = list(mocap_dir.glob("*.npz"))
    print(f"   ✓ Downloaded {len(npz_files)} motion capture files")

except ImportError:
    print("   ⚠️  huggingface_hub not available, skipping dataset download", file=sys.stderr)
    print("   You can download datasets manually from:", file=sys.stderr)
    print("   https://huggingface.co/datasets/robfiras/loco-mujoco-datasets", file=sys.stderr)
except Exception as e:
    print(f"   ⚠️  Error downloading datasets: {e}", file=sys.stderr)
    print("   You may need to run: huggingface-cli login", file=sys.stderr)
EOF

    echo "✓ Dataset download completed"
else
    # Count existing files
    FILE_COUNT=$(find "$MOCAP_DIR" -name "*.npz" | wc -l | tr -d ' ')
    echo "✓ Found existing mocap datasets ($FILE_COUNT files)"
fi

# Set up symlinks for OpenTrack
if [ -d "$OPENTRACK_REPO" ]; then
    echo ""
    echo "🔗 Setting up OpenTrack symlinks..."

    # Create symlink from OpenTrack/data/mocap to datasets
    OPENTRACK_DATA_DIR="$OPENTRACK_REPO/data/mocap"
    mkdir -p "$(dirname "$OPENTRACK_DATA_DIR")"

    if [ -L "$OPENTRACK_DATA_DIR" ]; then
        rm "$OPENTRACK_DATA_DIR"
    elif [ -d "$OPENTRACK_DATA_DIR" ]; then
        rm -rf "$OPENTRACK_DATA_DIR"
    fi

    ln -s "$DATASETS_DIR" "$OPENTRACK_DATA_DIR"
    echo "   ✓ $OPENTRACK_DATA_DIR -> $DATASETS_DIR"

    # Create symlink from OpenTrack/logs to models
    OPENTRACK_LOGS_DIR="$OPENTRACK_REPO/logs"

    if [ -L "$OPENTRACK_LOGS_DIR" ]; then
        rm "$OPENTRACK_LOGS_DIR"
    elif [ -d "$OPENTRACK_LOGS_DIR" ]; then
        # Move existing logs to MODELS_DIR
        if [ "$(ls -A "$OPENTRACK_LOGS_DIR" 2>/dev/null)" ]; then
            echo "   Moving existing logs to $MODELS_DIR..."
            mv "$OPENTRACK_LOGS_DIR"/* "$MODELS_DIR/" 2>/dev/null || true
        fi
        rm -rf "$OPENTRACK_LOGS_DIR"
    fi

    ln -s "$MODELS_DIR" "$OPENTRACK_LOGS_DIR"
    echo "   ✓ $OPENTRACK_LOGS_DIR -> $MODELS_DIR"
else
    echo "⚠️  OpenTrack repository not found at $OPENTRACK_REPO"
fi

echo ""
echo "======================================"
echo "✅ OpenTrack workspace initialized!"
echo ""
echo "   Datasets: $DATASETS_DIR"
echo "   Models:   $MODELS_DIR"
echo "   Videos:   $VIDEOS_DIR"
echo ""