#!/usr/bin/env python3
"""
Pre-download all models used in HuggingFace Enabling Sessions
Run this BEFORE your session to cache models locally
Models will save to ~/.cache/huggingface/hub
"""

import os
from transformers import pipeline, AutoTokenizer
from sentence_transformers import SentenceTransformer
import config

print("=" * 60)
print("🤗 HuggingFace Model Pre-Download Script")
print("=" * 60)

# Set HF cache directory (optional, for explicit control)
HF_HOME = os.path.expanduser("~/.cache/huggingface")
os.makedirs(HF_HOME, exist_ok=True)
print(f"\n📁 Cache location: {HF_HOME}")

models_to_download = [
    ("Sentiment Analysis", config.SENTIMENT_MODEL, "sentiment"),
    ("NER", config.NER_MODEL, "ner"),
    ("Question Answering", config.QA_MODEL, "qa"),
    ("Summarization", config.SUMMARIZATION_MODEL, "summarization"),
    ("Semantic Similarity", config.EMBEDDINGS_MODEL, "embeddings"),
]

print(f"\n📥 Starting download of {len(models_to_download)} models...\n")

# Download pipelines
for task_name, model_id, task_type in models_to_download[:4]:
    try:
        print(f"⏳ Downloading {task_name} ({model_id})...")
        if task_type == "ner":
            pipeline("ner", model=model_id)
        elif task_type == "qa":
            pipeline("question-answering", model=model_id)
        elif task_type == "summarization":
            pipeline("summarization", model=model_id)
        else:
            pipeline("sentiment-analysis", model=model_id)
        print(f"✅ {task_name} downloaded successfully\n")
    except Exception as e:
        print(f"❌ Error downloading {task_name}: {str(e)}\n")

# Download Sentence-BERT
try:
    print(f"⏳ Downloading Semantic Similarity ({config.EMBEDDINGS_MODEL})...")
    SentenceTransformer(config.EMBEDDINGS_MODEL)
    print(f"✅ Semantic Similarity downloaded successfully\n")
except Exception as e:
    print(f"❌ Error downloading Semantic Similarity: {str(e)}\n")

# Download tokenizer
try:
    print(f"⏳ Downloading Tokenizer ({config.SENTIMENT_MODEL})...")
    AutoTokenizer.from_pretrained(config.SENTIMENT_MODEL)
    print(f"✅ Tokenizer downloaded successfully\n")
except Exception as e:
    print(f"❌ Error downloading Tokenizer: {str(e)}\n")

print("=" * 60)
print("✅ Model pre-download complete!")
print("=" * 60)
print("\n📝 Notes:")
print("- All models are cached in ~/.cache/huggingface/hub")
print("- Models will be used instantly in Spaces demos")
print("- Total size: ~2-3 GB (may take 10-20 minutes)")
print("\n🚀 Ready for your session!")