#!/bin/bash
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 🚀 GaasHub: Unified Remote Execution (Unix V3 Hybrid)
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

START_TIME=$(date +%s)
SCUDA_DIR="$HOME/.scuda"
STATUS_FILE="$SCUDA_DIR/.connected"
CONFIG_FILE="$SCUDA_DIR/config"

# MacOS optimization
export COPYFILE_DISABLE=1
export _COPYFILE_DISABLE=1

# Colors
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; 
BLUE='\033[0;34m'; CYAN='\033[0;36m'; NC='\033[0m'

# 1. 🛑 CHECK CONNECTION
if [ ! -f "$STATUS_FILE" ]; then
    echo -e "${RED}❌ GaasHub Error: No connection metadata found at $SCUDA_DIR${NC}"
    echo -e "${YELLOW}Please open the GaasHub Desktop App and click 'Connect' to initialize.${NC}"
    exit 101
fi
source "$CONFIG_FILE"
SCUDA_HOST=$(echo "$SCUDA_SERVER" | cut -d':' -f1)
SCUDA_PORT=$(echo "$SCUDA_SERVER" | cut -d':' -f2)
SSH_PORT=${SSH_PORT:-$SCUDA_PORT}
SSH_KEY="$SCUDA_DIR/.gh_ident"
if [ ! -f "$SSH_KEY" ]; then SSH_KEY="$SCUDA_DIR/vast_ssh_key"; fi

# Optimized SSH Opts (Reduced Handshaking)
SSH_OPTS="-T -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -o LogLevel=ERROR -i $SSH_KEY"

# 2. 🔍 PARSING
TARGET="$1"
shift
REMOTE_ARGS="$@"

if [ -z "$TARGET" ]; then
    echo "Usage: gaashub <folder_or_script_or_command> [args]"
    exit 1
fi

# 3. 🧠 SMART RESOLUTION: Command vs Project
IS_LOCAL=0
if [ -e "$TARGET" ]; then
    IS_LOCAL=1
fi

if [ "$IS_LOCAL" -eq 1 ]; then
    # PROJECT MODE (Mac Mini Sync)
    if [ -d "$TARGET" ]; then
        PROJECT_DIR=$(cd "$TARGET" && pwd)
        ENTRY_FILE="main.py"
        [ ! -f "$PROJECT_DIR/main.py" ] && [ -f "$PROJECT_DIR/main.cu" ] && ENTRY_FILE="main.cu"
    else
        PROJECT_DIR=$(cd "$(dirname "$TARGET")" && pwd)
        ENTRY_FILE=$(basename "$TARGET")
    fi
    PROJECT_NAME=$(basename "$PROJECT_DIR")
    PARENT_DIR=$(dirname "$PROJECT_DIR")
    
    SESSION_ID="gh_${PROJECT_NAME}_$$"
    REMOTE_WORK_DIR="/tmp/$SESSION_ID"
    REMOTE_TAR="/tmp/${SESSION_ID}.tar.gz"
    EXCLUDES=(--exclude="__pycache__" --exclude=".git" --exclude=".venv" --exclude="node_modules" --exclude="runs" --exclude=".scuda")

    if [[ "$ENTRY_FILE" == *.py ]]; then RUN_CMD="python3 -u '$ENTRY_FILE' $REMOTE_ARGS"
    elif [[ "$ENTRY_FILE" == *.cu ]]; then RUN_CMD="nvcc -o m '$ENTRY_FILE' && ./m $REMOTE_ARGS"
    else RUN_CMD="bash '$ENTRY_FILE' $REMOTE_ARGS"
    fi

    REMOTE_SCRIPT="trap 'kill 0' EXIT; mkdir -p \"$REMOTE_WORK_DIR\" && cat > \"$REMOTE_TAR\" && tar -xzf \"$REMOTE_TAR\" -C \"$REMOTE_WORK_DIR\" && cd \"$REMOTE_WORK_DIR/$PROJECT_NAME\" && ([ -f 'requirements.txt' ] && pip install -r requirements.txt >/dev/null 2>&1 || true); if [ -f ~/.gaashub/scheduler.py ]; then python3 ~/.gaashub/scheduler.py \"$RUN_CMD\"; else $RUN_CMD; fi"

    echo -e "${YELLOW}📦 Syncing Workspace: ${NC}${CYAN}$PROJECT_NAME${NC}"
    echo -e "${GREEN}🚀 Running $ENTRY_FILE on Remote GPU...${NC}"
    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"

    tar "${EXCLUDES[@]}" -czf - -C "$PARENT_DIR" "$PROJECT_NAME" | ssh -p $SSH_PORT $SSH_OPTS root@$SCUDA_HOST "$REMOTE_SCRIPT"
    EXIT_CODE=$?
    
    # DOWNLOAD RESULTS
    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${YELLOW}📥 Fetching results...${NC}"
    LOCAL_RESULT_TAR="/tmp/res_${SESSION_ID}.tar.gz"
    ssh -p $SSH_PORT $SSH_OPTS root@$SCUDA_HOST "cd $REMOTE_WORK_DIR && tar ${EXCLUDES[@]} -czf - $PROJECT_NAME" > "$LOCAL_RESULT_TAR" 2>/dev/null
    if [ -f "$LOCAL_RESULT_TAR" ]; then
        tar -xzf "$LOCAL_RESULT_TAR" -C "$PARENT_DIR" 2>/dev/null
    fi
    
    # CLEANUP
    ssh -p $SSH_PORT $SSH_OPTS root@$SCUDA_HOST "rm -f $REMOTE_TAR; rm -rf $REMOTE_WORK_DIR" >/dev/null 2>&1
    rm -f "$LOCAL_RESULT_TAR" 2>/dev/null
else
    # COMMAND MODE (Instant Execution)
    echo -e "${GREEN}⚡ Executing System Command: ${NC}${CYAN}$TARGET $REMOTE_ARGS${NC}"
    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    ssh -p $SSH_PORT $SSH_OPTS root@$SCUDA_HOST "$TARGET $REMOTE_ARGS"
    EXIT_CODE=$?
fi

END_TIME=$(date +%s)
ELAPSED=$((END_TIME - START_TIME))
echo -e "${GREEN}✨ Finished. [Total: ${ELAPSED}s]${NC}"
exit $EXIT_CODE
