#!/usr/bin/env python3
import argparse
import hashlib
import inspect
import json
import os
import shutil
from pathlib import Path
import sys


def main():
    p=argparse.ArgumentParser()
    p.add_argument('--config', required=True); p.add_argument('--job', required=True)
    a=p.parse_args()
    from config import load_config
    cfg=load_config(a.config)
    runtime=Path(cfg['_paths']['runtime'])/'python'
    os.environ['PYTHONPATH']=str(runtime)
    sys.path.insert(0,str(runtime))
    import exoplasim
    from results import prepare_results
    job=json.loads(Path(a.job).read_text())
    spec=job['spec']; run=spec['run']; model=spec['model']
    work=Path(job['directory'])/'runtime'
    instance=exoplasim.Model(workdir=str(work),modelname='job_'+job['id'],resolution=model['resolution'],layers=model['layers'],ncpus=model['ncpus'])
    params=dict(spec.get('parameters',{}))
    map_paths=job.get('map_snapshots',{})
    resource_maps=[]
    steps=run.get('steps')
    params['runsteps']=steps
    if steps is not None:
        cadence=params.get('snapshots',inspect.signature(instance.configure).parameters['snapshots'].default)
        params['snapshots']=min(cadence,steps)
        instance.cfgpostprocessor(times=None,timeaverage=False)
    surface=spec.get('surface')
    if surface:
        from terrain import prepare_surface
        surface_config=dict(cfg); surface_config['parameters']=params
        maps=prepare_surface(surface,model,surface_config,str(work/'surface'))
        params['landmap']=maps['landmap']; params['topomap']=maps['topomap']
    if map_paths:
        from terrain import write_map_sra, _reference_gravity, _RESOLUTIONS
        gravity=params.get('gravity')
        if gravity is None: gravity=_reference_gravity(cfg)
        gravity=float(gravity)
        from mapcatalog import FIELDS
        if surface:
            landmask=map_paths.get('landmask') or maps['landmap']
        else:
            landmask=map_paths.get('landmask') or work/f"N{_RESOLUTIONS[model['resolution']]:03d}_surf_{FIELDS['landmask']['native_code']:04d}.sra"
        native_inputs=Path(job['directory'])/'inputs'/'native'
        native_inputs.mkdir(parents=True,exist_ok=True)
        for field, snapshot in map_paths.items():
            contract=FIELDS[field]
            native_path=write_map_sra(snapshot,model['resolution'],field,native_inputs/f'{field}.sra',gravity,landmask)
            if contract['native_parameter'] == 'resources':
                resource_maps.append((native_path,contract['native_code']))
            else:
                params[contract['native_parameter']]=native_path
    instance.configure(**params)
    # configure removes SRA files when height or land maps are supplied.
    for native_path,code in resource_maps:
        destination=work/f"N{_RESOLUTIONS[model['resolution']]:03d}_surf_{code:04d}.sra"
        shutil.copyfile(native_path,destination)
        instance.resources.append(str(destination))
    if steps is not None:
        for key in ('N_RUN_YEARS','N_RUN_MONTHS','N_RUN_DAYS'):
            instance._edit_namelist('plasim_namelist',key,'0')
    instance.exportcfg()
    jobroot=Path(job['directory'])
    for cycle in range(int(run['cycles'])):
        instance.run(years=1,crashifbroken=True,clean=False)
        output=work/f'MOST.{cycle:05d}.npz'
        if not output.is_file():
            candidates=sorted(work.glob('MOST.*.npz'))
            if not candidates: raise FileNotFoundError('Model produced no MOST NPZ output')
            output=candidates[-1]
        dataset_id=hashlib.sha256(f"{job['id']}:{cycle+1}".encode()).hexdigest()[:32]
        dest=jobroot/'artifacts'/'datasets'/dataset_id
        result=prepare_results(output,dest,dataset_id,config=cfg)
        budget=int(cfg.get('max_result_bytes',1024**3))
        used=sum(p.stat().st_size for p in (jobroot/'artifacts').rglob('*') if p.is_file())
        if used > budget: raise ValueError('Result output exceeds configured storage budget.')
        print(json.dumps({'cycle':cycle+1,'cycles':run['cycles'],'dataset':dataset_id,'output':str(output)}),flush=True)

if __name__=='__main__': main()