#!/usr/bin/env python3
"""Synthetic NetCDF test of the RegCM extraction chain.
Only sigma2p is substituted with an already pressure-level fixture.
All CDO/NCO commands are real, guarded, and all files are retained.
"""
import sys
sys.dont_write_bytecode=True
import json
from pathlib import Path
import numpy as np
from netCDF4 import Dataset
import rcms
cfg=json.loads((rcms.HERE/'config.json').read_text())
root=rcms.HERE/'test-results'/('regcm-post-'+rcms.stamp()); root.mkdir(parents=True)
raw=root/'raw'; raw.mkdir()
print('Synthetic RegCM outputs:',root,flush=True)
specs=[('ATM',6,'ua va wa ta hus rh hgt mslp ts'),('SRF',3,'pr prc uas vas ts huss hurs psl evspsbl zmla hfss hfls'),('STS',24,'pr prmax tas tasmax tasmin'),('RAD',6,'rlut prw')]
for family,step,names in specs:
    p=raw/('d01_'+family+'.2026082100.nc')
    with Dataset(str(p),'w',format='NETCDF3_CLASSIC') as d:
        for name,n in [('time',None),('lat',2),('lon',2),('plev',9)]: d.createDimension(name,n)
        t=d.createVariable('time','f8',('time',)); t.units='hours since 2026-08-21 00:00:00'; t.calendar='gregorian'; t[:]=list(range(0,25,step))
        lat=d.createVariable('lat','f8',('lat',)); lat.units='degrees_north'; lat[:]=[15,16]
        lon=d.createVariable('lon','f8',('lon',)); lon.units='degrees_east'; lon[:]=[105,106]
        z=d.createVariable('plev','f8',('plev',)); z.units='hPa'; z.axis='Z'; z.positive='down'; z[:]=[1000,925,850,700,600,500,300,200,100]
        for name in names.split():
            dims=('time','plev','lat','lon') if family=='ATM' and name not in ('mslp','ts') else ('time','lat','lon')
            v=d.createVariable(name,'f4',dims,fill_value=-9999.); v.units='1'
            shape=tuple(len(d.dimensions[k]) for k in dims)
            v[:]=np.arange(np.prod(shape)).reshape(shape)*0.1+1
class FixtureRunner(rcms.Runner):
    def run(self,args,cwd=None):
        if Path(str(args[0])).name=='sigma2pCLM45':
            src=Path(args[1]); dest=Path(cwd)/(src.stem+'_pressure.nc')
            rcms.copy_new(src,dest)
            rcms.put(Path(cwd)/'FIXTURE_SIGMA2P.json',{'note':'Synthetic source already on pressure levels; executable not tested'})
            return ''
        return super().run(args,cwd)
work=root/'post';work.mkdir()
r=FixtureRunner(work,cfg['utilities'])
try:
    products=rcms.regcm_post({'config':cfg,'cycle':'2026082100','end':'2026082200'},raw,work,r)
    count=len(list(products.rglob('*.nc')))
    assert count==28,count
    rcms.put(root/'SUCCESS.json',{'files':count,'sigma2p':'fixture only'})
    print('RegCM extraction smoke SUCCESS:',count,'products',flush=True)
except Exception as exc:
    rcms.put(root/'FAILED.json',{'error':str(exc)})
    raise