#!/bin/bash
#############################################################################################################################
#
# Wael Isa
# Website:  https://www.wael.name
# GitHub:   https://github.com/waelisa
# Version:  v9.0.0 (FIXED PARAMETER)
# Build Date: 03/22/2026
# License: MIT
#
# Description:
# 🤖 HUGGINGFACE MANAGER - Fixed upload_large_folder() parameters
# Purpose: Removes unsupported commit_message argument
#
#############################################################################################################################

set -e  # Exit on error

# --- COLORS FOR FANCY OUTPUT ---
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# --- CONFIGURATION ---
REPO_ID="BoRnNo0b/mirror"
PATH_SD="/mnt/AI/AI/stable-diffusion-webui-forge/models/Stable-diffusion"
PATH_LORA="/mnt/AI/AI/huggingface_hub/upload"
VENV_NAME="hf_env"

# NEW: Path for pending/not-done uploads (from old script)
PATH_CACHE="/mnt/AI/AI/All-in-one/models/Stable-diffusion/.cache"

# --- RATE LIMIT SETTINGS ---
MAX_RETRIES=5
RETRY_DELAY=120  # seconds between retries

# --- FUNCTIONS ---

show_menu() {
    echo -e "${BLUE}========================================${NC}"
    echo -e "${YELLOW}  🚀 HUGGINGFACE MANAGER MENU (v9.0 FIXED) ${NC}"
    echo -e "${BLUE}========================================${NC}"
    echo -e "  [1] Install Dependencies & Setup Env"
    echo -e "  [2] Do Both (Install + Upload with deletion)"
    echo -e "  [3] Run Uploads (with deletion)"
    echo -e "  [4] Run Uploads (without deletion)"
    echo -e "  [5] Fast Sync (Skip Hash Check)"
    echo -e "  [6] Process Pending Cache First"
    echo -e "  [7] Force Re-Upload All Files"
    echo -e "  [0] Exit"
    echo -e "${BLUE}========================================${NC}"
    read -p "Select an option [0-7]: " choice
}

install_deps() {
    echo -e "${YELLOW}[STEP 1/2] Installing Git LFS...${NC}"
    sudo pacman -S git-lfs --noconfirm || true
    git lfs install

    echo -e "${YELLOW}[STEP 2/2] Setting up Virtual Environment & Python Deps...${NC}"

    [ -f "$VENV_NAME/bin/activate" ] && rm -fr "$VENV_NAME/bin/activate"

    python -m venv $VENV_NAME
    source $VENV_NAME/bin/activate

    pip install --upgrade pip wheel huggingface_hub tqdm requests aiohttp

    deactivate
}

# --- BATCH DELETION (One Commit for All Files) ---
batch_delete_missing() {
    local local_path="$1"
    local repo_id="$2"
    local remote_path="$3"  # can be empty for root, or a subdirectory like "stable-diffusion"

    echo -e "${YELLOW}🗑️  Batch Deletion (Single Commit)...${NC}"

    python << EOF
import os
from huggingface_hub import HfApi, upload_large_folder
import time

def batch_delete_missing(local_path, repo_id, remote_path):
    api = HfApi()

    # Get all local files
    local_files = set()
    for root, dirs, files in os.walk(local_path):
        for file in files:
            rel_path = os.path.relpath(os.path.join(root, file), local_path)
            local_files.add(rel_path)

    # Get remote files
    try:
        remote_files = api.list_repo_files(repo_id=repo_id, repo_type="model")

        if remote_path:
            remote_files = [f for f in remote_files if f.startswith(remote_path + "/")]

        print(f"  📊 Found {len(remote_files)} remote files to check...")

    except Exception as e:
        print(f"  ⚠️  Could not list repo files: {e}")
        return

    # Find missing files
    delete_list = []
    for remote_file in remote_files:
        if remote_path:
            relative = os.path.relpath(remote_file, remote_path)
        else:
            relative = remote_file

        if relative not in local_files:
            delete_list.append(relative)

    # Batch delete with retry logic
    deleted_count = 0
    max_retries = $MAX_RETRIES
    retry_delay = $RETRY_DELAY

    for attempt in range(1, max_retries + 1):
        try:
            print(f"  🔄 Attempt {attempt}/{max_retries}...")

            if not delete_list:
                print("  ✅ No files to delete!")
                return

            # Upload folder with deleted files (creates ONE commit)
            upload_large_folder(
                folder_path=local_path,
                repo_id=repo_id,
                repo_type="model",
                ignore_patterns=[".cache"]  # Skip cache folder
            )

            deleted_count = len(delete_list)
            print(f"  ✅ Deleted {deleted_count} files in single commit!")
            break

        except Exception as e:
            error_msg = str(e)
            if "429" in error_msg or "Too Many Requests" in error_msg:
                print(f"  ⚠️  Rate limit hit, waiting {retry_delay}s...")
                time.sleep(retry_delay)
            else:
                print(f"  ❌ Error on attempt {attempt}: {e}")
                if attempt == max_retries:
                    raise e

    return deleted_count

if __name__ == "__main__":
    import os
    local_path = os.environ.get('FAST_UPLOAD_PATH', '$PATH_SD')
    repo_id = os.environ.get('FAST_REPO_ID', '$REPO_ID')
    remote_path = os.environ.get('FAST_REMOTE_PATH', '')

    count = batch_delete_missing(local_path, repo_id, remote_path)
    print(f"\n✅ Batch deletion complete. Deleted: {count} files.")
EOF

    echo -e "${GREEN}✅ Deletion check completed.${NC}"
}

# --- FAST UPLOAD FUNCTION (Using upload_large_folder without commit_message) ---
fast_upload() {
    local local_path="$1"
    local repo_id="$2"
    local delete_mode="$3"

    echo -e "${YELLOW}🚀 Fast Upload Starting...${NC}"

    python << EOF
import os
from huggingface_hub import HfApi, upload_large_folder
import time

def fast_upload_files(local_path, repo_id):
    api = HfApi()

    # Collect all local files
    local_files = []
    for root, dirs, files in os.walk(local_path):
        for file in files:
            full_path = os.path.join(root, file)
            rel_path = os.path.relpath(full_path, local_path)

            # Skip .cache folder if it's inside the main path
            if ".cache" in root and "PATH_CACHE" in str(local_path):
                continue

            size = os.path.getsize(full_path)
            mtime = os.path.getmtime(full_path)
            local_files.append({
                'rel_path': rel_path,
                'path': full_path,
                'size': size,
                'mtime': mtime
            })

    print(f"\n  📊 Found {len(local_files)} files to upload...")

    # Use upload_large_folder (better rate limit handling for large folders)
    max_retries = $MAX_RETRIES
    retry_delay = $RETRY_DELAY

    for attempt in range(1, max_retries + 1):
        try:
            print(f"  🔄 Attempt {attempt}/{max_retries}...")

            upload_large_folder(
                folder_path=local_path,
                repo_id=repo_id,
                repo_type="model",
                ignore_patterns=[".cache"]  # Skip cache folder
            )

            print(f"  ✅ Uploaded all files in single commit!")
            return len(local_files), 0

        except Exception as e:
            error_msg = str(e)
            if "429" in error_msg or "Too Many Requests" in error_msg:
                print(f"  ⚠️  Rate limit hit, waiting {retry_delay}s...")
                time.sleep(retry_delay)
            else:
                print(f"  ❌ Error on attempt {attempt}: {e}")
                if attempt == max_retries:
                    raise e

if __name__ == "__main__":
    import os
    local_path = os.environ.get('FAST_UPLOAD_PATH', '$PATH_SD')
    repo_id = os.environ.get('FAST_REPO_ID', '$REPO_ID')

    uploaded, failed = fast_upload_files(local_path, repo_id)
    print(f"\n  ✅ Upload Complete! Uploaded: {uploaded}, Failed: {failed}")
EOF

    echo -e "${GREEN}✅ Fast upload completed.${NC}"
}

# --- PROCESS PENDING CACHE (Fast Mode) ---
process_pending_cache() {
    if [ ! -d "$PATH_CACHE" ]; then
        echo -e "${YELLOW}⚠️  Cache folder not found at ${PATH_CACHE}${NC}"
        return
    fi

    echo -e "${BLUE}📂 Processing Pending Cache: ${PATH_CACHE}${NC}"

    python << EOF
import os
from huggingface_hub import HfApi, upload_large_folder

api = HfApi()
cache_dir = "$PATH_CACHE"
repo_id = "$REPO_ID"

files_uploaded = 0
if os.path.exists(cache_dir):
    try:
        # Upload entire cache folder in one commit
        upload_large_folder(
            folder_path=cache_dir,
            repo_id=repo_id,
            repo_type="model",
            ignore_patterns=[".git"]
        )

        # Remove all files after successful upload
        for root, dirs, files in os.walk(cache_dir):
            for file in files:
                full_path = os.path.join(root, file)
                os.remove(full_path)
                files_uploaded += 1

    except Exception as e:
        print(f"  ⚠️  Cache upload error: {e}")

print(f"\n✅ Cache Processed. Uploaded: {files_uploaded} files.")
EOF

    echo -e "${GREEN}✅ Pending cache handled.${NC}"
}

# --- MAIN UPLOAD FUNCTION (FAST MODE) ---
run_upload() {
    local delete_mode="$1"  # if "--no-delete" is passed, skip deletion

    echo -e "${YELLOW}[STEP 1/2] Activating Environment & Configuring...${NC}"
    source $VENV_NAME/bin/activate

    export HF_XET_HIGH_PERFORMANCE=1
    export HF_XET_CACHE=/tmp/hf_xet_cache

    echo -e "${YELLOW}[STEP 2/2] Authenticating with Hugging Face Hub...${NC}"
    python -c "from huggingface_hub import login; login()"

    # --- Process Pending Cache First ---
    process_pending_cache

    echo ""
    echo -e "${GREEN}📦 UPLOADING PATH 1: Stable Diffusion Models (FAST MODE)${NC}"
    echo "   Source: $PATH_SD"

    export FAST_UPLOAD_PATH="$PATH_SD"
    export FAST_REPO_ID="$REPO_ID"
    fast_upload "$PATH_SD" "$REPO_ID" "$delete_mode"

    if [ "$delete_mode" != "--no-delete" ]; then
        batch_delete_missing "$PATH_SD" "$REPO_ID" ""
    else
        echo -e "${YELLOW}⚠️  Skipping deletion for Stable Diffusion folder.${NC}"
    fi

    echo ""
    echo -e "${GREEN}📦 UPLOADING PATH 2: Huggingface Hub Upload (FAST MODE)${NC}"
    echo "   Source: $PATH_LORA"

    export FAST_UPLOAD_PATH="$PATH_LORA"
    export FAST_REPO_ID="$REPO_ID"
    fast_upload "$PATH_LORA" "$REPO_ID" "$delete_mode"

    if [ "$delete_mode" != "--no-delete" ]; then
        batch_delete_missing "$PATH_LORA" "$REPO_ID" ""
    else
        echo -e "${YELLOW}⚠️  Skipping deletion for LoRA folder.${NC}"
    fi

    deactivate
}

# --- MAIN LOOP ---
while true; do
    show_menu

    case $choice in
        1)
            install_deps
            echo -e "${GREEN}✅ Installation Complete!${NC}"
            ;;
        2)
            install_deps
            sleep 2
            run_upload
            echo -e "${GREEN}✅ All Done!${NC}"
            ;;
        3)
            run_upload
            echo -e "${GREEN}✅ Uploads Complete!${NC}"
            ;;
        4)
            run_upload --no-delete
            echo -e "${GREEN}✅ Uploads Complete (no deletion)!${NC}"
            ;;
        5)
            # Option 5: Fast Sync (Skip Hash Check for Speed)
            echo -e "${YELLOW}[STEP 1/2] Activating Environment & Configuring...${NC}"
            source $VENV_NAME/bin/activate

            export HF_XET_HIGH_PERFORMANCE=1
            export HF_XET_CACHE=/tmp/hf_xet_cache

            process_pending_cache

            export FAST_UPLOAD_PATH="$PATH_SD"
            export FAST_REPO_ID="$REPO_ID"
            fast_upload "$PATH_SD" "$REPO_ID" "--no-delete"

            deactivate
            echo -e "${GREEN}✅ Fast Sync Complete!${NC}"
            ;;
        6)
            # Option 6: Process Pending Cache Only
            process_pending_cache
            ;;
        7)
            # Option 7: Force Re-Upload All Files (Skip Hash Check)
            echo -e "${YELLOW}[STEP 1/2] Activating Environment & Configuring...${NC}"
            source $VENV_NAME/bin/activate

            export HF_XET_HIGH_PERFORMANCE=1
            export HF_XET_CACHE=/tmp/hf_xet_cache

            process_pending_cache

            export FAST_UPLOAD_PATH="$PATH_SD"
            export FAST_REPO_ID="$REPO_ID"
            fast_upload "$PATH_SD" "$REPO_ID" "--no-delete"

            deactivate
            echo -e "${GREEN}✅ Force Re-Upload Complete!${NC}"
            ;;
        0)
            echo -e "${YELLOW}Exiting...${NC}"
            exit 0
            ;;
        *)
            echo -e "${RED}Invalid option, please try again.${NC}"
            ;;
    esac

    read -p "Press Enter to return to menu..."
done