"""
Download full ERA5 dataset (1960-2026) for all 27 stations.
Downloads year by year with separate requests for instantaneous vs accumulated variables.

NOTE: ERA5 separates instantaneous (temperature, wind) from accumulated
(precipitation, evaporation, radiation) variables. They must be requested separately.
"""
import cdsapi
import pandas as pd
import xarray as xr
import zipfile
import shutil
from pathlib import Path

# Load station coordinates
stations = pd.read_csv('datasets/station_list.csv')
lat_min = stations['lat'].min() - 0.1
lat_max = stations['lat'].max() + 0.1
lon_min = stations['lon'].min() - 0.1
lon_max = stations['lon'].max() + 0.1

print("="*80)
print("ERA5 FULL DOWNLOAD (1960-2026)")
print("="*80)
print(f"Study area: [{lat_min:.2f}, {lon_min:.2f}] to [{lat_max:.2f}, {lon_max:.2f}]")
print(f"Years: 1960-2026 (67 years)")
print(f"Output: datasets/safran/")
print("="*80)

# Split variables by type (instantaneous vs accumulated)
instantaneous_vars = [
    '2m_temperature',               # T_Q
    '10m_u_component_of_wind',      # FF_Q
    '10m_v_component_of_wind',      # FF_Q
]

accumulated_vars = [
    'total_precipitation',                 # PRELIQ_Q
    'potential_evaporation',               # ETP_Q
    'surface_solar_radiation_downwards',   # DLI_Q / SSI_Q
    'snowfall',                            # PRENEI_Q
    'runoff',                              # RUNC_Q
]

def extract_netcdf(zip_file):
    """Extract NetCDF from ZIP archive."""
    if zipfile.is_zipfile(zip_file):
        with zipfile.ZipFile(zip_file, 'r') as zip_ref:
            # Find all .nc files in the zip
            nc_files = [f for f in zip_ref.namelist() if f.endswith('.nc')]
            extracted_files = []
            for nc_file in nc_files:
                temp_file = zip_file.parent / f'temp_{nc_file}'
                with zip_ref.open(nc_file) as source:
                    with open(temp_file, 'wb') as target:
                        shutil.copyfileobj(source, target)
                extracted_files.append(temp_file)
        zip_file.unlink()
        return extracted_files
    return [zip_file]

def merge_datasets(instant_file, accum_file, output_file):
    """Merge instantaneous and accumulated variable datasets."""
    ds_instant = xr.open_dataset(instant_file)
    ds_accum = xr.open_dataset(accum_file)

    # Merge datasets
    ds_merged = xr.merge([ds_instant, ds_accum])
    ds_merged.to_netcdf(output_file)

    ds_instant.close()
    ds_accum.close()
    instant_file.unlink()
    accum_file.unlink()

c = cdsapi.Client()
output_dir = Path('datasets/safran')
output_dir.mkdir(exist_ok=True)

# Download year by year (1960-2026)
years = list(range(1960, 2027))  # 1960 to 2026 inclusive
failed_years = []

for year in years:
    output_file = output_dir / f'era5_{year}.nc'

    # Skip if already exists and contains all 8 variables
    if output_file.exists():
        try:
            ds = xr.open_dataset(output_file)
            if len(ds.data_vars) >= 8:
                print(f"✓ {year} - Already downloaded with all variables, skipping")
                ds.close()
                continue
            else:
                print(f"⚠ {year} - Incomplete ({len(ds.data_vars)} vars), re-downloading")
                ds.close()
                output_file.unlink()
        except:
            output_file.unlink()

    print(f"\nDownloading {year}...")

    instant_file = output_dir / f'era5_{year}_instant.nc'
    accum_file = output_dir / f'era5_{year}_accum.nc'

    try:
        # Download instantaneous variables
        print(f"  → Instantaneous variables...")
        c.retrieve(
            'reanalysis-era5-single-levels',
            {
                'product_type': 'reanalysis',
                'format': 'netcdf',
                'variable': instantaneous_vars,
                'year': str(year),
                'month': [f'{m:02d}' for m in range(1, 13)],
                'day': [f'{d:02d}' for d in range(1, 32)],
                'time': '12:00',
                'area': [lat_max, lon_min, lat_min, lon_max],
            },
            str(instant_file)
        )

        # Extract if ZIP
        extracted = extract_netcdf(instant_file)
        if len(extracted) > 0:
            instant_file = extracted[0]

        # Download accumulated variables
        print(f"  → Accumulated variables...")
        c.retrieve(
            'reanalysis-era5-single-levels',
            {
                'product_type': 'reanalysis',
                'format': 'netcdf',
                'variable': accumulated_vars,
                'year': str(year),
                'month': [f'{m:02d}' for m in range(1, 13)],
                'day': [f'{d:02d}' for d in range(1, 32)],
                'time': '12:00',
                'area': [lat_max, lon_min, lat_min, lon_max],
            },
            str(accum_file)
        )

        # Extract if ZIP
        extracted = extract_netcdf(accum_file)
        if len(extracted) > 0:
            accum_file = extracted[0]

        # Merge both datasets
        print(f"  → Merging datasets...")
        merge_datasets(instant_file, accum_file, output_file)

        size_kb = output_file.stat().st_size / 1024
        print(f"✓ {year} - Downloaded ({size_kb:.1f} KB)")

    except Exception as e:
        print(f"✗ {year} - Failed: {e}")
        failed_years.append(year)
        # Cleanup partial files
        for f in [instant_file, accum_file]:
            if f.exists():
                f.unlink()
        continue

print("\n" + "="*80)
print("DOWNLOAD COMPLETE")
print("="*80)
print(f"Successful: {len(years) - len(failed_years)}/{len(years)} years")

if failed_years:
    print(f"\nFailed years: {failed_years}")
    print("You can re-run this script to retry failed downloads")
else:
    print("\n✓ All years downloaded successfully!")

print(f"\nFiles saved to: {output_dir}")
print("Total files:", len(list(output_dir.glob('era5_*.nc'))))