"""
Extract NetCDF files from ERA5 ZIP archives.
CDS API downloads NetCDF data as compressed ZIP files.
"""
import zipfile
from pathlib import Path
import shutil

safran_dir = Path('datasets/safran')

print("=" * 80)
print("EXTRACTING ERA5 NETCDF FILES FROM ZIP ARCHIVES")
print("=" * 80)

# Find all .nc files that are actually ZIP archives
nc_files = list(safran_dir.glob('era5_*.nc'))
extracted = 0
skipped = 0

for nc_file in sorted(nc_files):
    # Check if it's a ZIP file
    try:
        with zipfile.ZipFile(nc_file, 'r') as zip_ref:
            # It's a ZIP file - extract it
            zip_contents = zip_ref.namelist()

            # Find the actual NetCDF file inside (usually data.nc or similar)
            actual_nc = [f for f in zip_contents if f.endswith('.nc')][0]

            # Extract to temporary location
            temp_file = safran_dir / f'temp_{nc_file.name}'
            with zip_ref.open(actual_nc) as source:
                with open(temp_file, 'wb') as target:
                    shutil.copyfileobj(source, target)

            # Replace the ZIP with the extracted NetCDF
            nc_file.unlink()
            temp_file.rename(nc_file)

            size_kb = nc_file.stat().st_size / 1024
            print(f"✓ Extracted {nc_file.name} ({size_kb:.1f} KB)")
            extracted += 1

    except zipfile.BadZipFile:
        # Not a ZIP file - it's already a proper NetCDF, skip
        size_kb = nc_file.stat().st_size / 1024
        print(f"  {nc_file.name} - Already extracted ({size_kb:.1f} KB)")
        skipped += 1
    except Exception as e:
        print(f"✗ Failed to extract {nc_file.name}: {e}")

print("\n" + "=" * 80)
print("EXTRACTION COMPLETE")
print("=" * 80)
print(f"Extracted: {extracted} files")
print(f"Already extracted: {skipped} files")
print(f"Total NetCDF files: {len(nc_files)}")