src/utils/model_downloader.py

"""
Model Downloader - Downloads AI models from Hugging Face Hub
Automatically caches models locally after first download
FULLY PORTABLE - Works on any device with any project path
"""

from huggingface_hub import hf_hub_download
from pathlib import Path
import os
import sys
import shutil

# Detect PROJECT_ROOT dynamically
def get_project_root():
    """
    Find project root by looking for config/ directory
    Works regardless of where app.py is located
    """
    current_path = Path(__file__).resolve()  # Full path to this file
    
    # Go up from src/utils/model_downloader.py to project root
    for parent in current_path.parents:
        if (parent / 'config').exists() and (parent / 'webapp').exists():
            return parent
    
    # Fallback: assume parent of src/
    return current_path.parent.parent.parent

PROJECT_ROOT = get_project_root()
REPO_ID = "itsluckysharma01/NETRA-Models"
CACHE_DIR = PROJECT_ROOT / 'ai_models'  # Models cached in project root

print(f"\n🔍 [Model Downloader] PROJECT_ROOT detected: {PROJECT_ROOT}")
print(f"🔍 [Model Downloader] CACHE_DIR: {CACHE_DIR}\n")

def download_model(filename):
    """
    Download model from Hugging Face Hub with automatic path handling
    
    Args:
        filename: Model file path (e.g., 'ai_models/activity_recognition/violence_model.h5')
    
    Returns:
        str: Path to downloaded/cached model (absolute path)
    """
    try:
        # Ensure cache directory exists
        CACHE_DIR.mkdir(parents=True, exist_ok=True)
        
        # Check if model already exists in flat structure
        local_path = CACHE_DIR / filename
        if local_path.exists():
            print(f"✅ Model cached: {filename}")
            return str(local_path)
        
        # Download from Hugging Face Hub (goes to HF cache)
        print(f"📥 Downloading: {filename}")
        downloaded_path = hf_hub_download(
            repo_id=REPO_ID,
            filename=filename,
            cache_dir=str(CACHE_DIR),
            local_files_only=False
        )
        
        # Copy from HF cache structure to flat ai_models/ structure
        src_path = Path(downloaded_path)
        
        # Create destination directory
        local_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Copy file to flat structure
        shutil.copy2(src_path, local_path)
        print(f"✅ Downloaded and cached: {filename}")
        return str(local_path)
        
    except Exception as e:
        print(f"❌ Error downloading {filename}: {e}")
        return None


def ensure_model_exists(filename):
    """
    Ensure a model exists locally, download if necessary
    
    Args:
        filename: Model file path
    
    Returns:
        bool: True if model exists or was downloaded successfully
    """
    local_path = CACHE_DIR / filename
    
    # Already exists
    if local_path.exists():
        return True
    
    # Try to download
    result = download_model(filename)
    return result is not None


def setup_all_models():
    """Download all required models on startup"""
    models = [
        "ai_models/activity_recognition/violence_model.h5",
        "ai_models/object_detection/yolov8n.pt",
        "ai_models/pose_detection/yolo11n-pose.pt",
        "ai_models/weapon_detection/best.pt",
        "ai_models/analysis_models/binarycnn200.h5",
        "ai_models/analysis_models/CNN93.h5",
        "ai_models/analysis_models/CustomCNN.h5",
        "ai_models/analysis_models/fight_detection_model.h5",
    ]
    
    print("\n" + "=" * 60)
    print("📥 SETTING UP AI MODELS FROM HUGGING FACE HUB")
    print("=" * 60)
    print(f"🔍 PROJECT_ROOT: {PROJECT_ROOT}")
    print(f"🔍 CACHE_DIR: {CACHE_DIR}")
    print(f"🔍 Cache exists: {CACHE_DIR.exists()}")
    print("=" * 60)
    
    downloaded = 0
    cached = 0
    failed = 0
    
    for model in models:
        local_path = CACHE_DIR / model
        
        if local_path.exists():
            print(f"✅ Cached: {model}")
            cached += 1
        else:
            try:
                result = download_model(model)
                if result:
                    downloaded += 1
                else:
                    failed += 1
            except Exception as e:
                print(f"⚠️  Warning: Could not load {model}")
                failed += 1
    
    print("\n" + "=" * 60)
    print(f"✅ Setup Complete: {downloaded} downloaded, {cached} cached, {failed} warnings")
    print(f"📁 Models should be at: {CACHE_DIR}")
    print("=" * 60 + "\n")