"""One-request stdlib SSH agent. No server, listening socket, or credentials."""
import base64, contextlib, datetime, hashlib, json, os, re, shlex, subprocess, sys, tempfile, urllib.request, uuid
from pathlib import Path
REPO = 'will-taggart/AFFECTOR-ARC'

def atomic(path, obj):
    path = Path(path); path.parent.mkdir(parents=True, exist_ok=True)
    fd, name = tempfile.mkstemp(dir=path.parent)
    try:
        with os.fdopen(fd, 'w') as f:
            json.dump(obj, f, ensure_ascii=False, allow_nan=False); f.flush(); os.fsync(f.fileno())
        os.replace(name, path)
    finally:
        if os.path.exists(name): os.unlink(name)

def command(args, **kw):
    p = subprocess.run(args, capture_output=True, text=True, **kw)
    if p.returncode: raise RuntimeError((p.stderr or p.stdout or 'Command failed').strip())
    return p.stdout.strip()

def token(value, label):
    if not isinstance(value, str) or not re.fullmatch(r'[A-Za-z0-9_.,:+@/-]+', value) or value.startswith('-'):
        raise ValueError('Invalid '+label)
    return value

def rootpath(value):
    if not isinstance(value, str) or not value.strip(): raise ValueError('Choose a remote workspace')
    p = Path(value).expanduser()
    if not p.is_absolute() or p == Path('/') or len(p.parts)<3: raise ValueError('Use an absolute path or ~/folder for the workspace')
    return p.resolve()

def jobdir(root, identity):
    uuid.UUID(identity); p=root/'jobs'/identity
    if not p.is_dir(): raise ValueError('Job record not found')
    return p

def fetch(url):
    with urllib.request.urlopen(urllib.request.Request(url,headers={'User-Agent':'AFFECTOR-ARC-Desktop/1.0'}), timeout=120) as r: return r.read()

def current(root):
    p=root/'current'
    if not p.exists(): raise ValueError('Install/update AFFECTOR-ARC first')
    release=p.resolve()
    if release.parent != (root/'releases').resolve(): raise ValueError('Unexpected release path')
    return release

def envpath(root,release):
    digest=hashlib.sha256((release/'setup.sh').read_bytes()).hexdigest()[:16]
    return root/'environments'/digest

def scheduler(q, root, kind, payload=None):
    release=current(root); env=envpath(root,release)
    if kind!='setup' and not (env/'.affector-ready').exists(): raise ValueError('Run GPU setup successfully before submitting a transformation')
    identity=str(uuid.uuid4()); d=root/'jobs'/identity; d.mkdir(parents=True)
    account=token(q['account'],'allocation'); partition=token(q['partition'],'partition')
    wall=q.get('walltime','01:00:00')
    if not re.fullmatch(r'(?:\d+-)?\d{1,3}:\d{2}:\d{2}',wall): raise ValueError('Use HH:MM:SS or D-HH:MM:SS for wall time')
    cpus=int(q.get('cpus',8)); memory=int(q.get('memory',32))
    if not 1<=cpus<=128 or not 1<=memory<=1024: raise ValueError('Invalid CPU or memory request')
    cache=rootpath(q.get('cache',str(root/'model-cache')))
    lines=['#!/bin/bash','set -euo pipefail','umask 077','module load Miniconda3',f'export HF_HOME={shlex.quote(str(cache))}',f'cd {shlex.quote(str(release))}']
    if kind=='setup':
        # Reuse a completed environment, but never treat a partial setup as ready.
        lines += [f'ENV={shlex.quote(str(env))}', 'mkdir -p "$(dirname "$ENV")"', 'exec 9>"$ENV.lock"', 'flock 9', 'if [ ! -f "$ENV/.affector-ready" ]; then',
          'if [ ! -x "$ENV/bin/python" ]; then conda create -y --override-channels -c conda-forge -p "$ENV" python=3.12 pip; fi',
          '"$ENV/bin/python" -m pip install torch --index-url https://download.pytorch.org/whl/cu126',
          '"$ENV/bin/python" -m pip install "transformers==5.22.2" accelerate huggingface_hub safetensors',
          '"$ENV/bin/python" -m unittest discover -s tests',
          '"$ENV/bin/python" -m pip freeze > "$ENV/environment-resolved.txt"',
          '"$ENV/bin/python" -c \'import torch; assert torch.cuda.is_available(), "CUDA unavailable"; print(torch.__version__, torch.cuda.get_device_name(0))\'',
          'touch "$ENV/.affector-ready"', 'fi',
          '"$ENV/bin/python" -m unittest discover -s tests',
          '"$ENV/bin/python" -c \'import torch; assert torch.cuda.is_available(), "CUDA unavailable"; print(torch.__version__, torch.cuda.get_device_name(0))\'']
    else:
        text=q.get('text','')
        if not isinstance(text,str) or not text.strip(): raise ValueError('Input text is empty')
        if len(text.encode())>20_000_000: raise ValueError('Input exceeds 20 MB')
        (d/'input.txt').write_text(text,encoding='utf-8')
        geometry=q.get('geometry')
        if not isinstance(geometry,dict): raise ValueError('Missing shape')
        # Validate without importing torch. The pinned source supplies the actual schema.
        import importlib.util
        spec=importlib.util.spec_from_file_location('arc_geometry',release/'engine/geometry.py'); m=importlib.util.module_from_spec(spec); spec.loader.exec_module(m)
        geometry=m.import_atlas(geometry,geometry.get('name','Imported shape')); m.validate(geometry)
        atomic(d/'shape.json',geometry)
        args=[str(env/'bin/python'),'-u','run_arc.py','--mode',kind,'--text-file',str(d/'input.txt'),'--geometry',str(d/'shape.json'),'--output',str(d/'report.json')]
        style=q.get('style','Restrained')
        if style not in ['Minimal','Restrained','Balanced','Creative']: raise ValueError('Invalid rewrite style')
        args += ['--style',style]
        if kind=='rewrite':
            for key,lo,hi in [('temperature',0,2),('top_p',0.00001,1),('top_k',0,100000),('max_tokens',32,2048),('time_limit',1,3600),('seed',0,2147483647)]:
                val=q.get(key,{'temperature':0,'top_p':.95,'top_k':40,'max_tokens':512,'time_limit':120,'seed':0}[key])
                if type(val) not in (float,int) or not lo<=val<=hi: raise ValueError('Invalid '+key)
                if key in ['top_k','max_tokens','seed'] and int(val)!=val: raise ValueError(key+' must be an integer')
                args += ['--'+key.replace('_','-'),str(val)]
        lines += [shlex.join(args)]
    script='\n'.join(lines)+'\n'; (d/'job.sh').write_text(script)
    jid=command(['sbatch','--parsable','--account='+account,'--partition='+partition,'--nodes=1','--ntasks=1','--cpus-per-task='+str(cpus),'--gres=gpu:1','--mem='+str(memory)+'G','--time='+wall,'--job-name=affector-'+kind,'--output='+str(d/'job.log'),'--error='+str(d/'job.log'),str(d/'job.sh')]).split(';')[0]
    if not jid.isdigit(): raise RuntimeError('Unexpected Slurm job ID: '+jid)
    record={'id':identity,'job_id':jid,'kind':kind,'created_at':datetime.datetime.now(datetime.timezone.utc).isoformat(),'release':release.name,'directory':str(d),'account':account,'partition':partition}
    atomic(d/'job.json',record); return record


def response_text(text):
    if not isinstance(text,str):return ''
    text=re.sub(r'<think>[\s\S]*?</think>\s*','',text,flags=re.I)
    if re.search(r'<think>',text,re.I):return ''
    label=r'(?:transformed(?:\s+(?:text|passage|document))?|rewritten(?:\s+(?:text|passage|document))?|rewrite|revised(?:\s+(?:text|passage|document))?|final(?:\s+(?:answer|text))?|attempt|output)'
    pattern=r'(?im)^\s*(?:\#{1,6}\s*)?(?:here is (?:the |a )?)?(?:\*\*)?'+label+r'(?:\*\*)?\s*:\s*(?:\*\*)?[^\S\r\n]*'
    matches=list(re.finditer(pattern,text))
    if matches:
        text=text[matches[-1].end():]
        suffix=re.search(r'(?im)^\s*(?:\#{1,6}\s*)?(?:\*\*)?(?:original(?:\s+(?:text|target|passage))?|source(?:\s+text)?|prompt|target|instructions?|explanation|notes?)(?:\*\*)?\s*:',text)
        if suffix:text=text[:suffix.start()]
    elif re.search(r'(?im)^\s*(?:\#{1,6}\s*)?(?:\*\*)?(?:original(?:\s+(?:text|target|passage))?|source\s+text|prompt|target)(?:\*\*)?\s*:',text):return ''
    text=text.strip()
    if text.startswith('```') and text.endswith('```'):
        lines=text.splitlines();text='\n'.join(lines[1:-1]).strip()
    if re.match(r"(?is)^(?:the prompt (?:asks|requests|requires|instructs)|this is (?:a |an )?prohibited attempt|i (?:cannot|can't|am unable to) (?:rewrite|comply with|fulfill this request))\b",text):return ''
    return text

def transformed_text(report):
    records=report.get('records')
    if not isinstance(records,dict):return response_text(report.get('text',''))
    spans=report.get('spans',[]);source=report.get('input','');parts=[];previous=None
    for index in sorted(int(k) for k in records if str(k).isdigit()):
        record=records[str(index)];text=record.get('text','')
        if record.get('original_retained') or not isinstance(text,str) or not text.strip():continue
        if previous is not None:
            gap='\n\n'
            if index==previous+1 and index<len(spans) and len(spans[previous])==2 and len(spans[index])==2:
                a,b=spans[previous][1],spans[index][0]
                if 0<=a<=b<=len(source) and not source[a:b].strip():gap=source[a:b]
            parts.append(gap)
        clean=response_text(text)
        if not clean:
            if previous is not None:parts.pop()
            continue
        parts.append(clean);previous=index
    return ''.join(parts)

def dispatch(q):
    op=q['op']; root=rootpath(q['workspace'])
    if op=='probe':
        return {'host':command(['hostname']),'home':str(Path.home()),'workspace':str(root),'slurm':command(['sbatch','--version']),'installed':(root/'current').exists(),'revision':current(root).name if (root/'current').exists() else None,'ready':(envpath(root,current(root))/'.affector-ready').exists() if (root/'current').exists() else False}
    root.mkdir(parents=True,exist_ok=True)
    if op=='update':
        meta=json.loads(fetch('https://huggingface.co/api/models/'+REPO));rev=meta['sha']
        if not re.fullmatch('[0-9a-f]{40}',rev):raise RuntimeError('Invalid HF revision')
        release=root/'releases'/rev
        if not (release/'.snapshot-complete').exists():
            tree=json.loads(fetch(f'https://huggingface.co/api/models/{REPO}/tree/{rev}?recursive=true&limit=1000'))
            selected=[x for x in tree if x['type']=='file' and (x['path'].startswith(('engine/','tests/')) or x['path'] in ['run_arc.py','document_run.py','setup.sh','smoke.sbatch','README.md','README.txt','ATLAS_LICENSE','ENGINE_MANIFEST.json'])]
            if not selected or not {'run_arc.py','setup.sh'}.issubset({x['path'] for x in selected}):raise RuntimeError('Incomplete HF snapshot')
            release.mkdir(parents=True,exist_ok=True)
            for entry in selected:
                rel=Path(entry['path'])
                if rel.is_absolute() or '..' in rel.parts:raise RuntimeError('Unsafe snapshot path')
                dest=release/rel;dest.parent.mkdir(parents=True,exist_ok=True)
                data=fetch(f'https://huggingface.co/{REPO}/resolve/{rev}/{rel.as_posix()}')
                if data.startswith(b'version https://git-lfs.github.com/spec/'):raise RuntimeError('Unresolved LFS pointer')
                dest.write_bytes(data)
            (release/'.snapshot-complete').write_text(rev)
        temp=root/('current-'+str(uuid.uuid4()));temp.symlink_to(release,target_is_directory=True);os.replace(temp,root/'current')
        return {'revision':rev,'path':str(release),'ready':(envpath(root,release)/'.affector-ready').exists()}
    if op in ['setup','rewrite','smoke']:return scheduler(q,root,op)
    if op=='list':
        result=[]
        for p in (root/'jobs').glob('*/job.json'):
            with contextlib.suppress(Exception):result.append(json.loads(p.read_text()))
        return {'jobs':sorted(result,key=lambda x:x['created_at'],reverse=True)}
    if op in ['status','cancel','download']:
        d=jobdir(root,q['id']);record=json.loads((d/'job.json').read_text());jid=record['job_id']
        if not jid.isdigit():raise ValueError('Invalid job ID')
        if op=='cancel':command(['scancel',jid]);return {'message':'Cancellation requested','job_id':jid}
        if op=='download':
            files={}
            for name in ['input.txt','shape.json','report.json','report.txt','job.log','job.json','job.sh']:
                p=d/name
                if p.exists():files[name]=base64.b64encode(p.read_bytes()).decode()
            return {'files':files}
        try: queued=command(['squeue','-h','-j',jid,'-o','%T|%R'])
        except RuntimeError: queued=''
        state=queued.split('|',1)[0] if queued else ''
        if not state:
            try:state=command(['sacct','-n','-X','-j',jid,'--format=State','--parsable2']).splitlines()[0].split('|')[0].strip()
            except Exception:state='UNKNOWN'
        log=(d/'job.log').read_text(errors='replace')[-100000:] if (d/'job.log').exists() else ''
        data={'job':record,'state':state,'queue':queued,'log':log}
        p=d/'report.json'
        if p.exists():
            try:
                report=json.loads(p.read_text())
                # A lightweight status response; downloads retain the full basis.
                data['text']=report.get('text','')
                data['transformed_text']=transformed_text(report)
                preview=dict(report)
                if isinstance(preview.get('geometry'),dict):
                    preview['geometry']=dict(preview['geometry'],basis='Full calibration weights are in the downloaded report.')
                data['report']=preview
            except (ValueError,OSError):pass
        textfile=d/'report.txt'
        if not data.get('text') and textfile.exists():
            data['text']=textfile.read_text(errors='replace')
        return data
    raise ValueError('Unknown operation')

if __name__=='__main__':
    try:
        q=json.load(sys.stdin);print(json.dumps({'ok':True,'data':dispatch(q)},ensure_ascii=False,allow_nan=False))
    except Exception as e:
        print(json.dumps({'ok':False,'error':str(e)}));sys.exit(1)