"""Download recent FTP update files for testing"""
import ftplib, gzip, io, sys, time

HOST = "ftp.ncbi.nlm.nih.gov"
PATH = "/pubmed/updatefiles/"
FILES = ["pubmed26n1544.xml.gz", "pubmed26n1545.xml.gz", "pubmed26n1546.xml.gz"]

for name in FILES:
    print(f"Downloading {name}...", end=" ", flush=True)
    for attempt in range(3):
        try:
            ftp = ftplib.FTP(HOST, timeout=120)
            ftp.login()
            ftp.cwd(PATH)
            buf = io.BytesIO()
            ftp.retrbinary(f"RETR {name}", buf.write)
            ftp.quit()
            data = buf.getvalue()
            # verify
            gzip.decompress(data)
            with open(name, "wb") as f:
                f.write(data)
            print(f"OK ({len(data)/1024/1024:.1f}MB)")
            break
        except Exception as e:
            print(f"attempt {attempt+1} failed: {e}")
            time.sleep(3)
    else:
        print("FAILED")
        sys.exit(1)

print("All files downloaded and verified successfully")