import sys
import urllib.request

from fragment.protein_fragments import constants


def download_pdb_files(pdb_id, save_dir):
    filename = f'{pdb_id}.pdb'
    file_path = save_dir/filename
    if not file_path.exists():
        try:
            urllib.request.urlretrieve(f'https://files.rcsb.org/download/{filename}', file_path)
        except Exception as e:
            print(f'Failed to download {filename}: {e}')


def download(pdb_id):
    if not constants.RAW_DATA_DIR.exists():
        constants.RAW_DATA_DIR.mkdir(exist_ok=True)

    download_pdb_files(pdb_id, constants.RAW_DATA_DIR)


if __name__ == "__main__":
    assert len(sys.argv) == 2
    download(sys.argv[1])
