"""Finite, process-isolated boundary experiment; NOT the conditioning runner.

Requires external preregistration SHA and exclusive outputs. No automatic retry.
The old AP worker is invoked as an executable only, never imported as a neural
module into this coordinator. Every child exit and output inventory is checked.
"""
from __future__ import annotations

import argparse
import copy
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time
import traceback

import numpy as np

import boundary_plan as PLAN
import execution_authority as A
import isolated_execution as W
from prepare_execution import ROOT, MODEL, EXP, AP, LAB, save, save_npz, npz, verify, copy_exclusive

require,read,digest=A.require,A.read,A.digest
AP_WORKER=MODEL.parent/'conditioning-interface-v2/isolated_workers.py'


def registration(expected):
    file=EXP/'protocol/preregistered.json'
    require(digest(file)==expected,'External registration SHA changed')
    reg=read(file)
    require(reg['schema']=='conditioning-execution-registration-v1' and reg['engineering_only'] is True and
            reg['scientific_acquisition_authorized'] is False,'Wrong experiment authority')
    PLAN.validate_plan(reg['plan'])
    verify({**reg['source_sha256'],**reg['input_sha256'],**reg['prepared_file_sha256']})
    return reg


def tree(directory,*,exclude_build=False):
    result={}
    for p in sorted(Path(directory).rglob('*')):
        require(not p.is_symlink(),'Linked output forbidden')
        if exclude_build and p.relative_to(directory).parts[0]=='extension-build':continue
        if p.is_file():result[p.relative_to(directory).as_posix()]=dict(sha256=digest(p),bytes=p.stat().st_size)
        else:require(p.is_dir(),'Nonregular output forbidden')
    return result


class Workflow:
    def __init__(self,sha):
        self.started=time.perf_counter();self.sha=sha;self.reg=registration(sha)
        self.deadline=self.started+self.reg['budget']['workflow_timeout_seconds']
        self.parents=copy.deepcopy(self.reg['parent_records']);self.jobs={};self.boundaries={}
        self.completed=[];self.cycles=[];self.sampler=None
        self.h=W.ap_helpers()

    def admission(self):
        require(time.perf_counter()<self.deadline,'Whole-workflow wall limit reached')
        require(not any((EXP/n).exists() for n in ('audit/stopped.json','audit/execution-complete.json','manifest.json')),
                'Stopped/completed workflow cannot be retried')
        budget=self.reg['budget'];used=sum(p.stat().st_size for p in EXP.rglob('*') if p.is_file())
        require(used<=budget['output_cap_bytes'] and shutil.disk_usage(EXP).free>=
                budget['minimum_free_reserve_bytes']+max(0,budget['output_cap_bytes']-used),'Finite disk cap/reserve exceeded')

    def allpins(self):
        return {**self.reg['source_sha256'],**self.reg['input_sha256'],**self.reg['prepared_file_sha256'],
                str(EXP/'protocol/preregistered.json'):self.sha}

    def launch(self,name,job,*,ap=False):
        self.admission()
        loc={n:EXP/f'{sub}/{name}{suffix}' for n,sub,suffix in
             (('job','jobs','.json'),('receipt','receipts','.worker.json'),('process','receipts','.process.json'),
              ('stdout','logs','.stdout.log'),('stderr','logs','.stderr.log'))}
        require(not any(p.exists() for p in loc.values()),'Exclusive process record exists; no retry')
        require(job['receipt_file']==str(loc['receipt']),'Receipt location differs')
        save(loc['job'],job);jobsha=digest(loc['job'])
        script=AP_WORKER if ap else MODEL/'isolated_execution.py'
        command=[sys.executable,'-B',str(script)]
        if ap:command.append(job['kind'])
        command+=['--job',str(loc['job']),'--job-sha256',jobsha]
        start=time.perf_counter();started_utc=datetime.now(timezone.utc).isoformat()
        timeout=min(self.reg['budget']['child_timeout_seconds'],self.deadline-start)
        require(timeout>0,'No remaining workflow wall allowance')
        timed_out=False;exit_code=None
        with loc['stdout'].open('xb') as out,loc['stderr'].open('xb') as err:
            try:exit_code=subprocess.run(command,cwd=MODEL,stdout=out,stderr=err,timeout=timeout).returncode
            except subprocess.TimeoutExpired:timed_out=True
        exited=time.perf_counter()
        base=dict(schema='conditioning-execution-process-v1',name=name,kind=job['kind'],registration_sha256=self.sha,
            argv=command,job_sha256=jobsha,exit_code=exit_code,timed_out=timed_out,started_utc=started_utc,
            spawn_monotonic=start,exit_monotonic=exited,child_wall_seconds=exited-start,
            stdout_sha256=digest(loc['stdout']),stderr_sha256=digest(loc['stderr']))
        try:
            require(exit_code==0 and not timed_out and loc['receipt'].is_file(),'Child failed; preserve partial outputs')
            receipt=read(loc['receipt']);expected='phase4ap-isolated-worker-receipt-v1' if ap else 'conditioning-isolated-execution-receipt-v1'
            require(receipt['schema']==expected and receipt['complete'] is True and receipt['job_sha256']==jobsha and
                    receipt['registration_sha256']==self.sha and receipt['kind']==job['kind'] and
                    receipt['job_id']==job['job_id'] and receipt['worker_sha256']==job['worker_sha256'],
                    'Real worker receipt identity differs')
            expectedpins={**job['input_sha256'],str(loc['job']):jobsha,str(script):job['worker_sha256']}
            if not ap:
                expectedpins[str(AP_WORKER)]=W.AP_SHA
                expectedpins.update({str(MODEL/n):s for n,s in job['source_sha256'].items()})
            else:expectedpins.update({v['path']:v['sha256'] for v in job['sources'].values()})
            require(receipt['verified_input_sha256' if ap else 'input_sha256']==expectedpins,
                    'Child did not verify exact external source/input closure')
            actual={k:v['sha256'] for k,v in tree(Path(job['output_dir'])).items()}
            require(actual and actual==receipt['output_files_sha256'] and receipt['output_dir']==job['output_dir'],
                    'Actual output inventory differs from successful worker')
            require(digest(str(loc['receipt'])+'.attempt.json')==receipt['attempt_sha256'],'Attempt record changed')
            verify(expectedpins);self.admission()
            end=time.perf_counter()
            record=dict(base,complete=True,worker_receipt_sha256=digest(loc['receipt']),
                validation_end_monotonic=end,end_to_end_seconds=end-start,validation_wall_seconds=end-exited,
                output_bytes=sum(v['bytes'] for v in tree(Path(job['output_dir'])).values()),
                output_inventory_sha256=hashlib.sha256(json.dumps(actual,sort_keys=True).encode()).hexdigest(),
                timing_scope='before child spawn through actual exit and independent parent file verification')
            save(loc['process'],record)
        except BaseException as exc:
            if not loc['process'].exists():save(loc['process'],dict(base,complete=False,error_type=type(exc).__name__,error=str(exc)))
            raise
        self.jobs[name]=dict(loc={k:str(v) for k,v in loc.items()},job=job,process=record,receipt=receipt)
        print(json.dumps(dict(child=name,complete=True,end_to_end_seconds=record['end_to_end_seconds']),sort_keys=True),flush=True)
        return receipt

    def new_job(self,name,kind,authority,*,output=None,**extra):
        auth=read(authority);pins=self.allpins();pins[str(authority)]=digest(authority)
        if kind=='sample':pass
        else:pins.update(A.files_from_authority(auth))
        return dict(schema=W.SCHEMA,job_id=name,kind=kind,registration_sha256=self.sha,
            worker_sha256=self.reg['source_sha256'][str(MODEL/'isolated_execution.py')],
            source_sha256={n:self.reg['source_sha256'][str(MODEL/n)] for n in W.SOURCES[kind]},input_sha256=pins,
            output_dir=str(output or EXP/'runs'/name),receipt_file=str(EXP/'receipts'/f'{name}.worker.json'),
            budget={k:self.reg['budget'][k] for k in ('root','output_cap_bytes','minimum_free_reserve_bytes')},
            authority_file=str(authority),authority_sha256=digest(authority),**extra)

    def ap_job(self,name,kind,*,case=None,overlay=None,authority=None,output):
        pins=self.allpins()
        job=dict(schema='phase4ap-isolated-worker-job-v1',job_id=name,kind=kind,registration_sha256=self.sha,
            worker_sha256=W.AP_SHA,output_dir=str(output),receipt_file=str(EXP/'receipts'/f'{name}.worker.json'))
        if kind=='brain':
            for path in (case/'spec.json',case/'requests.npz',overlay['overlay_file'],overlay['overlay_manifest']):pins[str(path)]=digest(path)
            job.update(case_dir=str(case),build_root=str(EXP/'extension-build'),**overlay,diagnostic_mode='light' if name.startswith('performance') else 'reference',
                admission=dict(gpu_uuid=self.h.GPU,comfy_queue_urls=list(self.h.COMFY_URLS)),
                sources={n:dict(path=str(LAB/(n+'.py')),sha256=self.h.SOURCE_PINS[n]) for n in ('runtime_light','runtime_continuation')})
        else:
            value=read(authority);pins.update(A.files_from_authority(value));pins[str(authority)]=digest(authority)
            job.update(authority_file=str(authority),authority_sha256=digest(authority),
                sources={'conditioning_epoch':dict(path=str(LAB/'conditioning_epoch.py'),sha256=self.h.SOURCE_PINS['conditioning_epoch'])})
        job['input_sha256']=pins
        return job

    def cpu_verify(self,key,parent):
        name=key+'--cpu-verify';authority=Path(parent['epoch_authority']['path']);update=Path(parent['update_dir'])
        job=self.new_job(name,'validate_cpu',authority,update_dir=str(update),chain_sha256=digest(update/'chain-receipt.json'))
        # Original epoch authority doesn't include its future update files.
        job['input_sha256'].update({str(p):digest(p) for p in update.iterdir()})
        self.launch(name,job)
        parent['cpu_validation']=str(Path(job['output_dir'])/'validation.json')

    def materialize(self,call,parent=None):
        group=call['case_group'];case=EXP/'cases'/group
        spec=read(self.reg['case_templates'][group])
        if parent is not None:
            spec.update(parent_checkpoint_manifest_sha256=parent['checkpoint_sha256'],global_tick_offset=call['global_tick_offset'])
        if call['runner']=='runtime_light' and call['call_id'].startswith('performance'):
            request=Path(self.parents[self.reg['plan']['imported_parents']['D1']['role']]['epoch_authority']['path'])
            source=Path(read(request)['files']['requests']['path']);arrays=None
        elif call['runner']=='runtime_light':source=EXP/'inputs/static/mechanical-requests.npz';arrays=None
        elif call['runner']=='boundary_oracle':
            pp=self.parents[call['prefix_replay_parent']]
            arrays=npz(read(pp['epoch_authority']['path'])['files']['requests']['path']);source=None
        elif call['input_policy']=='externally_pinned_CS_probe':
            require(self.sampler is not None,'Evaluation input not actually sampled')
            source=Path(self.sampler['output_dir'])/f"{call['cue']}-requests.npz";arrays=None
        else:source=None;arrays=dict(tick=np.empty(0,np.int32),neuron_idx=np.empty(0,np.int32))
        if case.exists():
            require(read(case/'spec.json')==spec,'Shared case metadata differs')
            if source:require(digest(case/'requests.npz')==digest(source),'Shared saved inputs differ')
            else:require(all(npz(case/'requests.npz')[k].tobytes()==v.tobytes() for k,v in arrays.items()),'Shared request arrays differ')
        else:
            case.mkdir();save(case/'spec.json',spec)
            if source:copy_exclusive(source,case/'requests.npz')
            else:save_npz(case/'requests.npz',arrays)
        return case

    def sources(self,epoch):
        files=copy.deepcopy(epoch['files'])
        roles=('runtime_continuation_source','runtime_light_source','conditioning_update_source',
               'eligibility_bridge_source','phase4j_source','am_overlay_source')
        result={r:files[r] for r in roles}
        for role,name in (('conditioning_epoch_source','conditioning_epoch.py'),
                          ('conditioning_preflight_compare_source','conditioning_preflight_compare.py'),
                          ('engineering_compare_source','engineering_compare.py')):result[role]=A.pin(LAB/name)
        for role,name in (('updated_continuation_source','updated_continuation.py'),('cpu_validator_source','validate_update_cpu.py')):
            result[role]=A.pin(MODEL/name)
        return result

    def boundary(self,call,case,parent):
        transition=call['transition'];epoch=read(parent['epoch_authority']['path'])
        before=parent['before_overlay'] if transition=='epoch_update' else parent['after_overlay']
        path=EXP/'authorities'/f"{call['call_id']}.json"
        value=A.boundary_authority(boundary_id=call['call_id'],registration_sha256=self.sha,case=case,parent=parent['directory'],
            transition=transition,input_policy=call['input_policy'],mode=call['diagnostic_mode'],
            before_overlay=before,after_overlay=parent['after_overlay'],source_files=self.sources(epoch),frozen_inputs=epoch['frozen_inputs'],
            epoch=parent['epoch_authority']['path'] if transition=='epoch_update' else None,
            update_dir=parent['update_dir'] if transition=='epoch_update' else None,
            cpu_receipt=parent['cpu_validation'] if transition=='epoch_update' else None)
        save(path,value);self.boundaries[call['call_id']]=path
        return path

    def fresh(self,call):
        name=call['call_id'];case=self.materialize(call)
        performance=name.startswith('performance')
        prior=self.parents[self.reg['plan']['imported_parents']['D1']['role']]
        overlay=prior['before_overlay'] if performance else self.reg['initial_overlay']
        output=EXP/'runs'/name
        result=self.launch(name,self.ap_job(name,'brain',case=case,overlay=overlay,output=output),ap=True)
        auth=EXP/'authorities'/f'{name}.epoch.json'
        value=A.epoch_authority(read(prior['epoch_authority']['path']),case=case,checkpoint=output,overlay=overlay,
                                epoch_id=name,us_tick=900 if performance else 1199)
        save(auth,value);update=EXP/'updates'/name
        self.launch(name+'--original-E',self.ap_job(name+'--original-E','update',authority=auth,output=update),ap=True)
        parent=dict(directory=str(output),checkpoint_sha256=result['result']['checkpoint_manifest_sha256'],
            epoch_authority=A.pin(auth),update_dir=str(update),before_overlay=overlay,
            after_overlay=dict(overlay_file=str(update/'delta-after.npz'),overlay_manifest=str(update/'overlay-manifest.json')))
        self.cpu_verify(name,parent);self.parents[name]=parent
        return result

    def execute_call(self,call):
        name=call['call_id'];runner=call['runner']
        if runner=='runtime_light':return self.fresh(call)
        if runner=='fresh_probe':
            require(call.get('US') is False and call.get('plastic_update') is False and
                    call.get('feedback_to_idle') is False and call['parent'] is None and
                    call['overlay_after_role']=='AP.initial_overlay','Fresh probe scope differs')
            case=self.materialize(call)
            output=EXP/'runs'/name
            result=self.launch(name,self.ap_job(name,'brain',case=case,overlay=self.reg['initial_overlay'],output=output),ap=True)
            return result
        if runner=='boundary_oracle':
            ref='mechanical--updated--reference' if name.startswith('mechanical') else 'D1--updated--reference'
            authority=self.boundaries[ref];case=self.materialize(call)
            job=self.new_job(name,'oracle',authority,build_root=str(EXP/'extension-build'),oracle_case=str(case),
                oracle_spec_sha256=digest(case/'spec.json'),oracle_requests_sha256=digest(case/'requests.npz'))
            job['input_sha256'].update({str(case/n):digest(case/n) for n in ('spec.json','requests.npz')})
        else:
            parent=self.parents[call['parent']];case=self.materialize(call,parent)
            if runner=='runtime_continuation':
                authority=self.boundaries[name.split('--')[0]+'--updated--reference']
                job=self.new_job(name,'unchanged',authority,case_dir=str(case),build_root=str(EXP/'extension-build'),
                    diagnostic_mode='reference',checkpoint_dir=parent['directory'],checkpoint_sha256=parent['checkpoint_sha256'],**parent['before_overlay'])
            else:
                if call['input_policy']=='externally_pinned_CS_probe':
                    require(call.get('US') is False and call.get('plastic_update') is False and call.get('feedback_to_idle') is False,
                            'Probe cannot request US, updates, or feed idle')
                authority=self.boundary(call,case,parent)
                job=self.new_job(name,'updated',authority,build_root=str(EXP/'extension-build'),diagnostic_mode=call['diagnostic_mode'])
        result=self.launch(name,job)
        if runner=='updated_continuation':
            parent=self.parents[call['parent']]
            self.parents[name]={**copy.deepcopy(parent), 'directory':job['output_dir'],
                'checkpoint_sha256':result['result']['checkpoint_manifest_sha256'],
                'before_overlay':parent['after_overlay']}
        return result

    def run(self):
        attempt=EXP/'audit/execution-attempt.json'
        require(not attempt.exists(),'Workflow already attempted; no automatic retry')
        save(attempt,dict(registration_sha256=self.sha,started_utc=datetime.now(timezone.utc).isoformat(),
                          pid=os.getpid(),monotonic_origin=self.started,no_automatic_retry=True))
        for arm,descriptor in self.reg['plan']['imported_parents'].items():self.cpu_verify('AP-'+arm,self.parents[descriptor['role']])
        for call in self.reg['plan']['calls']:
            if call['input_policy']=='externally_pinned_CS_probe' and self.sampler is None:
                job=self.new_job('evaluation-inputs','sample',EXP/'protocol/preregistered.json')
                self.sampler=self.launch('evaluation-inputs',job)
            started=time.perf_counter();before=tree(EXP,exclude_build=True)
            build_before=tree(EXP/'extension-build') if (EXP/'extension-build').is_dir() else {}
            self.execute_call(call)
            after=tree(EXP,exclude_build=True)
            build_after=tree(EXP/'extension-build') if (EXP/'extension-build').is_dir() else {}
            end=time.perf_counter()
            self.admission()
            require(all(k in after and after[k]==v for k,v in before.items()),'Earlier output was modified')
            created={k:v for k,v in after.items() if k not in before}
            record=dict(call=call,complete=True,start_monotonic=started,end_monotonic=end,wall_seconds=end-started,
                output_inventory=created,output_bytes=sum(v['bytes'] for v in created.values()),
                mutable_build_inventory_before=build_before,mutable_build_inventory_after=build_after,
                mutable_build_bytes_before=sum(v['bytes'] for v in build_before.values()),
                mutable_build_bytes_after=sum(v['bytes'] for v in build_after.values()),
                inventory_scope='all new immutable artifacts; separate actual mutable build before/after inventory; shared prior setup excluded',
                source_registration_sha256=self.sha,learning_pass_claimed=False)
            save(EXP/'receipts'/f"{call['call_id']}.cycle.json",record)
            self.cycles.append(record);self.completed.append(call['call_id'])
            self.admission()
        require(len(self.completed)==self.reg['plan']['counts']['new_brain_calls'] and
                len(self.jobs)==self.reg['plan']['counts']['total_children'],'Actual finite call inventory differs')
        self.admission()
        save(EXP/'audit/execution-complete.json',dict(complete=True,registration_sha256=self.sha,
            completed_calls=self.completed,processes={k:v['loc']['process'] for k,v in self.jobs.items()},
            boundary_authorities={k:str(v) for k,v in self.boundaries.items()},parents=self.parents,
            wall_seconds=time.perf_counter()-self.started,numerical_comparison_complete=False,
            scientific_learning_pass_claimed=False,candidate_generated=False))


if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--registration-sha256',required=True)
    args=parser.parse_args()
    try:Workflow(args.registration_sha256).run()
    except BaseException as exc:
        stop=EXP/'audit/stopped.json'
        if stop.parent.is_dir() and not stop.exists():save(stop,dict(complete=False,error_type=type(exc).__name__,error=str(exc),
            registration_sha256=args.registration_sha256,retry_forbidden=True,traceback=traceback.format_exc()))
        raise