"""JSON worker for pinned moPPIt and PeptiVerse installations.

Run inside the scientific model's own environment with this package installed.
Artifacts are written to a persistent work directory supplied by the operator.
"""
import argparse,csv,hashlib,importlib.util,json,os,subprocess,sys
from pathlib import Path
from .schema import Candidate,Molecule,ToolResult,canonical

def run(mode,root,work,request,output):
    root=Path(root).resolve();work=Path(work).resolve();work.mkdir(parents=True,exist_ok=True)
    d=json.loads(Path(request).read_text());a=d['arguments'];state=d['state']
    if mode=='moppit':
        length=int(a.get('length',12));count=int(a.get('count',32))
        if length not in state['spec']['sequence_lengths'] or not 1<=count<=96:raise ValueError('generation limits')
        objectives=a.get('objectives',['Affinity','Solubility'])
        if set(objectives)-{'Affinity','Solubility','Hemolysis','Non-Fouling','Permeability','Half-Life','Motif'}:raise ValueError('unsupported native objective')
        weights=a.get('weights',[1/len(objectives)]*len(objectives))
        if len(weights)!=len(objectives) or any(float(x)<0 for x in weights) or sum(weights)<=0:raise ValueError('invalid weights')
        native_order=['Hemolysis','Non-Fouling','Solubility','Permeability','Half-Life','Affinity','Motif']
        weight_map=dict(zip(objectives,weights));objectives=[x for x in native_order if x in weight_map];weights=[weight_map[x] for x in objectives]
        if int(a.get('steps',100))!=100:raise ValueError('pinned native entrypoint fixes 100 integration steps')
        if state['spec']['peptide_format']!='linear':raise ValueError('this worker implements canonical linear moPPIt proposals')
        target=state['spec']['targets'][0]['molecule']['sequence']
        job=hashlib.sha256(canonical(d).encode()).hexdigest()[:16];csvpath=work/(job+'.csv')
        if csvpath.exists():raise ValueError('output exists; use a new episode or explicit cached result')
        argv=[sys.executable,'moppit.py','--target_protein',target,'--length',str(length),'--n_samples','1','--n_batches',str(count),
              '--T',str(int(a.get('steps',100))),'--objectives',*objectives,'--weights',*[str(float(x)) for x in weights],'--output_file',str(csvpath)]
        if 'Motif' in objectives:argv+=['--motifs',str(a['native_motif_indices'])]
        # Set all relevant random seeds before executing the native CLI.
        bootstrap='import runpy,sys,random,numpy as np,torch; s=int(sys.argv[1]); random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s); sys.argv=sys.argv[2:]; runpy.run_path(sys.argv[0],run_name="__main__")'
        seed=int(state['spec']['seed']);argv=[sys.executable,'-c',bootstrap,str(seed),*argv[1:]]
        with (work/(job+'.log')).open('w') as log:
            subprocess.run(argv,cwd=root,stdout=log,stderr=subprocess.STDOUT,check=True,timeout=7200)
        with csvpath.open() as f:rows=list(csv.DictReader(f))
        if not rows:raise ValueError('native generator returned no candidates')
        candidates=[Candidate(molecule=Molecule(sequence=r['Binder'].strip()),generator='ChatterjeeLab/moPPIt',revision='29de994bfb6c67890d0efca25d1cf431eea96999') for r in rows]
        result=ToolResult(candidates=candidates,artifacts={'native_scores':str(csvpath)},message='Native scores retained in CSV with original definitions; attach calibration before conservative ranking.')
    elif mode=='peptiverse':
        sys.path.insert(0,str(root));os.chdir(root)
        module_spec=importlib.util.spec_from_file_location('peppaverse_native',root/'inference.py');module=importlib.util.module_from_spec(module_spec);module_spec.loader.exec_module(module)
        predictor=module.PeptiVersePredictor(manifest_path=str(root/'best_models.txt'),classifier_weight_root=str(root),device=a.get('device','cuda'))
        raw={};target=state['spec']['targets'][0]['molecule']['sequence']
        for cid in a['candidate_ids']:
            mol=Molecule.model_validate(state['candidates'][cid]['molecule'])
            if mol.modifications or mol.bonds or mol.n_terminus!='free' or mol.c_terminus!='free':raise ValueError('use a chemistry-compatible native model for modified candidates')
            properties=a.get('properties',['solubility','hemolysis'])
            allowed={'solubility','hemolysis','nf','permeability_penetrance','halflife','affinity'}
            if set(properties)-allowed:raise ValueError('unsupported sequence endpoint')
            raw[cid]={}
            for prop in properties:
                value=predictor.predict_binding_affinity(col='wt',target_seq=target,binder_str=mol.sequence) if prop=='affinity' else predictor.predict_property(prop,col='wt',input_str=mol.sequence)
                raw[cid][prop]=value
        job=hashlib.sha256(canonical(d).encode()).hexdigest()[:16];p=work/(job+'_raw.json')
        def serial(x):
            if hasattr(x,'tolist'):return x.tolist()
            if hasattr(x,'item'):return x.item()
            raise TypeError(type(x).__name__)
        p.write_text(json.dumps(raw,default=serial,indent=2,allow_nan=False))
        result=ToolResult(artifacts={'peptiverse_native':str(p)},message='Native predictions preserved. Use endpoint-specific calibration before ranking.')
    else:raise ValueError(mode)
    Path(output).write_text(result.model_dump_json(indent=2))

def main():
    p=argparse.ArgumentParser();p.add_argument('--mode',choices=['moppit','peptiverse'],required=True);p.add_argument('--root',required=True);p.add_argument('--work',required=True);p.add_argument('request');p.add_argument('output');run(**vars(p.parse_args()))
if __name__=='__main__':main()