from huggingface_hub import hf_hub_download
import os

def download_assets():
    """Download necessary assets from Hugging Face Hub"""
    # Create directories if they don't exist
    os.makedirs('data', exist_ok=True)
    os.makedirs('vectorstore/db_faiss', exist_ok=True)
    
    # Dataset repository ID
    repo_id = "MoizK/mindmedic-assets"
    
    # Download PDF files
    pdf_files = [
        "71763-gale-encyclopedia-of-medicine.-vol.-1.-2nd-ed.pdf",
        "Depression-NIM-2024.pdf",
        "Depression-and-Other-Common-Mental-Disorders-Global-Health-Estimates.pdf",
        "Doing-What-Matters-in-Times-of-Stress.pdf",
        "Generalized-Anxiety-Disorder-When-Worry-Gets-Out-of-Control.pdf",
        "WHO-mhGAP-Intervention-Guide-v2.pdf",
        "social-anxiety-disorder-more-than-just-shyness.pdf"
    ]
    
    for pdf_file in pdf_files:
        try:
            hf_hub_download(
                repo_id=repo_id,
                filename=f"data/{pdf_file}",
                local_dir=".",
                local_dir_use_symlinks=False
            )
            print(f"Downloaded {pdf_file}")
        except Exception as e:
            print(f"Error downloading {pdf_file}: {e}")
    
    # Download FAISS index files
    index_files = ["index.faiss", "index.pkl"]
    for index_file in index_files:
        try:
            hf_hub_download(
                repo_id=repo_id,
                filename=f"vectorstore/db_faiss/{index_file}",
                local_dir=".",
                local_dir_use_symlinks=False
            )
            print(f"Downloaded {index_file}")
        except Exception as e:
            print(f"Error downloading {index_file}: {e}")

if __name__ == "__main__":
    download_assets()