import os
import shutil
from huggingface_hub import snapshot_download

# Set up paths
MODEL_REPO = "sentence-transformers/all-MiniLM-L6-v2"
MODEL_CACHE_DIR = "/app/model_cache"

print("⏳ Downloading the SentenceTransformer model...")
model_path = snapshot_download(repo_id=MODEL_REPO, cache_dir=MODEL_CACHE_DIR)
print("Model path now: ", model_path)

# Find the correct snapshot directory
snapshots_dir = os.path.join(model_path, "snapshots") # snapshot
if os.path.exists(snapshots_dir):
    snapshot_folders = os.listdir(snapshots_dir)
    if snapshot_folders:
        snapshot_dir = os.path.join(snapshots_dir, snapshot_folders[0])  # Get first snapshot folder
        print(f"📂 Moving model files from {snapshot_dir} to {MODEL_CACHE_DIR}...")
        # Move files from snapshot directory to model cache directory
        for filename in os.listdir(snapshot_dir):
            shutil.move(os.path.join(snapshot_dir, filename), os.path.join(MODEL_CACHE_DIR, filename))
        print(f"✅ Model extracted to {MODEL_CACHE_DIR}")
    else:
        print("❌ No snapshot folder found!")
else:
    print("❌ No snapshots directory found!")

# Verify structure
print("\n📂 LLM Model Structure:")
for root, dirs, files in os.walk(MODEL_CACHE_DIR):
    print(f"📁 {root}/")
    for file in files:
        print(f"  📄 {file}")