import urllib.request
import os
import SimpleITK as sitk

SAMPLES = [
    {
        "name": "sample_ct_chest.nii.gz",
        "url": "https://github.com/Slicer/SlicerTestingData/releases/download/SHA256/4507b664690840abb6cb9af2d919377ffc4ef75b167cb6fd0f747befdb12e38e"
    },
    {
        "name": "sample_ct_abdomen_panoramix.nii.gz",
        "url": "https://github.com/Slicer/SlicerTestingData/releases/download/SHA256/146af87511520c500a3706b7b2bfb545f40d5d04dd180be3a7a2c6940e447433"
    },
    {
        "name": "sample_ct_cardio.nii.gz",
        "url": "https://github.com/Slicer/SlicerTestingData/releases/download/SHA256/3b0d4eb1a7d8ebb0c5a89cc0504640f76a030b4e869e33ff34c564c3d3b88ad2"
    },
    {
        "name": "sample_ct_liver.nii.gz",
        "url": "https://github.com/Slicer/SlicerTestingData/releases/download/SHA256/e16eae0ae6fefa858c5c11e58f0f1bb81834d81b7102e021571056324ef6f37e"
    },
    {
        "name": "sample_ct_teeth.nii.gz",
        "url": "https://github.com/Slicer/SlicerTestingData/releases/download/SHA256/7bfa16945629c319a439f414cfb7edddd2a97ba97753e12eede3b56a0eb09968",
        "source_ext": ".gipl.gz"
    }
]

def download_and_convert():
    examples_dir = "examples"
    os.makedirs(examples_dir, exist_ok=True)
    
    for sample in SAMPLES:
        output_path = os.path.join(examples_dir, sample["name"])
        
        if "source_ext" in sample:
            temp_ext = sample["source_ext"]
        elif "gipl.gz" in sample["name"] or "gipl.gz" in sample["url"]:
            temp_ext = ".gipl.gz"
        elif "nii.gz" in sample["name"]:
            temp_ext = ".nii.gz"
        else:
            temp_ext = ".nrrd"
            
        print(f"DEBUG: sample={sample['name']}, source_ext={sample.get('source_ext')}, temp_ext={temp_ext}")
            
        temp_file = os.path.join(examples_dir, f"temp_download{temp_ext}")
        
        if os.path.exists(output_path):
            print(f"Skipping {sample['name']}, already exists.")
            continue
            
        print(f"Downloading {sample['name']} from {sample['url']}...")
        try:
            # Use a custom generic compatible header to avoid 403 blocks sometimes
            opener = urllib.request.build_opener()
            opener.addheaders = [('User-agent', 'Mozilla/5.0')]
            urllib.request.install_opener(opener)
            
            urllib.request.urlretrieve(sample['url'], temp_file)
            
            print(f"  Converting {temp_ext} to NIfTI...")
            img = sitk.ReadImage(temp_file)
            sitk.WriteImage(img, output_path)
            
            os.remove(temp_file)
            print(f"  ✓ Success: {output_path}")
            
        except Exception as e:
            print(f"  ✗ Failed: {e}")
            if os.path.exists(temp_file):
                os.remove(temp_file)

if __name__ == "__main__":
    download_and_convert()