#!/bin/bash
# Khudi AI v2 — On-start script (v3 - ALL FIXES APPLIED)
# All issues from v1 + v2.1 addressed

# Don't use 'set -e' - we want to handle errors gracefully
# instead of exiting on first failure

echo "============================================"
echo "🇵🇰 Khudi AI v2 — SAFE TRAINING (v3)"
echo "============================================"
echo "Start: $(date)"

# ============================================================
# FIX 1: Initialize IDLE_COUNT before health check
# ============================================================
IDLE_COUNT=0

# ============================================================
# SAFETY: 14-hour auto-shutdown
# ============================================================
( sleep 50400 && echo "⏰ 14 HOUR TIMEOUT" && pkill -9 -f "python train_v2" ) &
TIMEOUT_PID=$!

# ============================================================
# SAFETY: Health check (only after training starts)
# ============================================================
start_health_check() {
    while true; do
        sleep 300
        GPU_UTIL=$(nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits 2>/dev/null || echo "0")
        echo "[HEALTH] $(date) GPU: ${GPU_UTIL}% idle_count=$IDLE_COUNT"
        if [ "${GPU_UTIL:-0}" -lt 5 ]; then
            IDLE_COUNT=$((IDLE_COUNT + 1))
            if [ "$IDLE_COUNT" -gt 12 ]; then
                echo "[HEALTH] GPU idle 60+ min - EMERGENCY EXIT"
                # Save what we have
                if [ -d /workspace/khudi-v2-output/checkpoints ]; then
                    LATEST=$(ls /workspace/khudi-v2-output/checkpoints/ 2>/dev/null | grep checkpoint- | sort -V | tail -1)
                    if [ -n "$LATEST" ]; then
                        echo "Saving $LATEST to HF..."
                        HF_TOKEN="$HF_TOKEN_VALUE" python -c "
from huggingface_hub import HfApi
import os
api = HfApi()
api.upload_folder(
    folder_path='/workspace/khudi-v2-output/checkpoints/$LATEST',
    repo_id='ZaoKing/khudi-ai-v2',
    repo_type='model',
    commit_message='EMERGENCY: $LATEST - health check timeout'
)
print('Saved $LATEST to HF')
" 2>&1
                    fi
                fi
                pkill -9 -f "python train_v2" 2>/dev/null
                kill -9 $TIMEOUT_PID 2>/dev/null
                exit 2
            fi
        else
            IDLE_COUNT=0
        fi
    done
}

# ============================================================
# Step 1: Install packages
# ============================================================
echo "[1/8] Installing packages..."
pip install -q --no-cache-dir \
    "transformers==4.45.0" \
    "datasets==2.20.0" \
    "peft==0.11.0" \
    "trl==0.10.0" \
    "accelerate==0.34.0" \
    "bitsandbytes==0.43.3" \
    "huggingface_hub==0.25.0" \
    "sentencepiece" "protobuf" 2>&1 | tail -3 || echo "  ⚠️ Some packages may have failed"

# ============================================================
# Step 2: Setup directories
# ============================================================
echo "[2/8] Setting up dirs..."
mkdir -p /workspace
cd /workspace
mkdir -p khudi-v2-output/checkpoints khudi-v2-output/hf_cache

# ============================================================
# Step 3: Download train script from HF
# ============================================================
echo "[3/8] Downloading training scripts from HF..."
HF_TOKEN_VALUE="HF_TOKEN_PLACEHOLDER"

# FIX: Add retry logic
download_with_retry() {
    local url="$1"
    local output="$2"
    local max_retries=3
    local retry=0

    while [ $retry -lt $max_retries ]; do
        if wget -q --header="Authorization: Bearer $HF_TOKEN_VALUE" -O "$output" "$url" && [ -s "$output" ]; then
            return 0
        fi
        retry=$((retry + 1))
        echo "  Retry $retry/$max_retries for $output..."
        sleep 5
    done
    return 1
}

if ! download_with_retry "https://huggingface.co/ZaoKing/khudi-ai-v2/resolve/main/train_v2.py" "train_v2.py"; then
    echo "❌ FATAL: Could not download train_v2.py"
    exit 1
fi
echo "  ✅ train_v2.py downloaded: $(wc -l < train_v2.py) lines"

# ============================================================
# Step 4: Download v2 dataset (4 chunks)
# ============================================================
echo "[4/8] Downloading v2 dataset (4 chunks)..."
for i in 0 1 2 3; do
    echo "  Downloading chunk $i..."
    if ! download_with_retry "https://huggingface.co/ZaoKing/khudi-ai-v2/resolve/main/data/v2_chunk_$i.jsonl" "v2_chunk_$i.jsonl"; then
        echo "  ❌ FATAL: Could not download chunk $i"
        exit 1
    fi
done

# Combine chunks
cat v2_chunk_0.jsonl v2_chunk_1.jsonl v2_chunk_2.jsonl v2_chunk_3.jsonl > v2_data.jsonl
TOTAL=$(wc -l < v2_data.jsonl)
SIZE=$(ls -lh v2_data.jsonl | awk '{print $5}')
echo "  ✅ Combined: $TOTAL samples ($SIZE)"

# ============================================================
# Step 5: Quick data validation
# ============================================================
echo "[5/8] Validating data..."
HF_TOKEN="$HF_TOKEN_VALUE" python -c "
import json
count = 0
errors = 0
with open('/workspace/v2_data.jsonl') as f:
    for line in f:
        try:
            item = json.loads(line)
            if 'messages' not in item:
                errors += 1
                continue
            count += 1
        except:
            errors += 1
print(f'Valid: {count}, Errors: {errors}')
if count < 1000:
    print('FATAL: Not enough valid samples!')
    exit(1)
" || { echo "❌ Data validation failed"; exit 1; }

# ============================================================
# Step 6: Check for resume checkpoint
# ============================================================
echo "[6/8] Checking for resume checkpoints..."
if [ -d khudi-v2-output/checkpoints ] && [ "$(ls -A khudi-v2-output/checkpoints 2>/dev/null)" ]; then
    LATEST=$(ls khudi-v2-output/checkpoints/ | grep checkpoint- | sort -V | tail -1)
    echo "  ✅ Found checkpoint: $LATEST - will resume"
    RESUME_FLAG="--resume_from_checkpoint"
else
    echo "  No checkpoints - starting fresh"
    RESUME_FLAG=""
fi

# ============================================================
# Step 7: Start training
# ============================================================
echo "[7/8] Starting health check + training..."
start_health_check &
HEALTH_PID=$!
echo "  Health check PID: $HEALTH_PID"

# Export ALL required env vars
export HF_TOKEN="$HF_TOKEN_VALUE"
export HUGGINGFACE_HUB_TOKEN="$HF_TOKEN_VALUE"
export DATA_PATH="/workspace/v2_data.jsonl"
export OUTPUT_DIR="/workspace/khudi-v2-output"
export MODEL_NAME="Qwen/Qwen3.5-9B"

# FIX: Pass --epochs 1 explicitly
cd /workspace
python train_v2.py $RESUME_FLAG --epochs 1 2>&1

TRAIN_EXIT=$?
echo "Training exit: $TRAIN_EXIT"
echo "End: $(date)"

# Stop health check
kill $HEALTH_PID 2>/dev/null || true
kill $TIMEOUT_PID 2>/dev/null || true

# ============================================================
# Step 8: Final status
# ============================================================
echo "[8/8] Final status..."
if [ $TRAIN_EXIT -eq 0 ]; then
    echo "✅ Training completed"
else
    echo "⚠️ Training had issues - exit $TRAIN_EXIT"
fi

# Upload final log
HF_TOKEN="$HF_TOKEN_VALUE" python -c "
from huggingface_hub import HfApi
api = HfApi()
api.upload_file(
    path_or_fileobj='/workspace/train.log',
    path_in_repo='train_v2_$(date +%Y%m%d_%H%M%S).log',
    repo_id='ZaoKing/khudi-ai-v2',
    repo_type='model',
    commit_message='Training log v2'
)
print('Log uploaded')
" 2>&1 || echo "  Log upload failed"

# Stop instance
echo "Stopping instance..."
pkill -9 -f "python train_v2" 2>/dev/null || true
sleep 2
sudo shutdown -h now 2>/dev/null || shutdown -h now 2>/dev/null || true
echo "Script complete"