import os
import sys
from pathlib import Path
from huggingface_hub import snapshot_download

def main():
    if len(sys.argv) < 3:
        print("Usage: python download_model.py <repo_id> <local_dir>", file=sys.stderr)
        return 1
    
    repo_id = sys.argv[1]
    local_dir = Path(sys.argv[2])
    
    # Check if the local directory exists and contains files (excluding hidden ones like .gitattributes)
    has_files = False
    if local_dir.exists():
        for item in local_dir.iterdir():
            if item.is_file() and not item.name.startswith('.'):
                has_files = True
                break
            elif item.is_dir():
                has_files = True
                break
                
    if has_files:
        print(f"Model {repo_id} is already present at {local_dir}. Skipping download.")
        return 0
        
    # Fictional/Mock models check to prevent build failures
    if repo_id.startswith("CoExpressionLabs/coexpression-llm-global-3.3b"):
        print(f"Skipping download for fictional placeholder model: {repo_id}")
        return 0
        
    print(f"Model {repo_id} not found at {local_dir}. Downloading from Hugging Face...")
    token = os.environ.get("HF_TOKEN") or None
    
    local_dir.mkdir(parents=True, exist_ok=True)
    try:
        # Enable hf_transfer for fast downloads if available
        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
        import hf_transfer
    except ImportError:
        pass
        
    try:
        snapshot_download(
            repo_id=repo_id,
            local_dir=str(local_dir),
            local_dir_use_symlinks=False,
            token=token
        )
        print(f"Successfully downloaded {repo_id} to {local_dir}")
        return 0
    except Exception as e:
        print(f"Error downloading {repo_id}: {e}", file=sys.stderr)
        # For this hackathon, we allow downloads of other placeholder/fictional models to fail gracefully
        if "404" in str(e) or "gated" in str(e).lower():
            print(f"Warning: could not download {repo_id} (could be fictional or gated). Continuing build.")
            return 0
        return 1

if __name__ == "__main__":
    sys.exit(main())