"""Prepare the pinned release and one reusable model worker on every startup."""
from pathlib import Path
from contextlib import nullcontext
import atexit
import copy
import json
import os
import subprocess
import sys
import threading
import time
import traceback
import uuid

from hf_space import native_build, generation_progress
from hf_space.gpu_mode import enabled as zero_gpu_enabled

ROOT=Path(__file__).resolve().parents[1]
JOBS_ROOT=Path('/tmp/supratex-jobs')
MODEL_REVISION='1bc02eb969a8ed0f1418c2bbe1e409a8481f834c'
CODE_REVISION='e4dfbbde6a8d76e25deef5aba5d836f03994967d'
TEXTURE_FILTER_REVISION='cee4435652d0f23737bd28d5704e0fae9830b411'
CHECKPOINT='denoiser_step0030000.safetensors'
MODEL_ROOT=Path('/data/models')
STATE=Path('/tmp/supratex-release-preparation.json')
_lock=threading.Lock()
_worker=None
_foreground_worker=None
_foreground_started=False
_epoch=os.environ.get('SUPRATEX_STARTUP_EPOCH') or uuid.uuid4().hex
# Generated binding identities include resolved cache paths. Never reuse a
# previous process's prepare marker after a cache/transport upgrade.
READY=Path('/tmp/supratex-latest-ready')/_epoch
_started=False
_runtime_process=None
_runtime_directory=Path('/tmp/supratex-runtime')/_epoch
_generation_lock=threading.Lock()
_state_lock=threading.Lock()
_persist_lock=threading.Lock()
_snapshot={'stage':'pending','epoch':_epoch}
_foreground_snapshot={'stage':'pending','epoch':_epoch}
_staging_process=None
_runtime_monitor=None
_shutdown=threading.Event()
_startup_deadline=None
STARTUP_TOTAL_SECONDS=7200
STARTUP_ATTEMPTS=2
STARTUP_BACKOFF_SECONDS=15
NATIVE_WAIT_SECONDS=6000
STAGING_SECONDS=3600
STAGING_IDLE_SECONDS=300
RUNTIME_SECONDS=3600
RUNTIME_IDLE_SECONDS=900

def model_status():
    # Authentication and progress polling are pure reads. Only boot's unique
    # controller monitors children; readers cannot start/cancel/retry loading.
    with _state_lock:
        return copy.deepcopy(_snapshot)


def foreground_status():
    if not zero_gpu_enabled():
        return model_status()
    with _state_lock:
        return copy.deepcopy(_foreground_snapshot)


def _update_foreground(**changes):
    global _foreground_snapshot
    with _state_lock:
        _foreground_snapshot = _foreground_snapshot | changes | {'updated_unix': time.time()}
    print('[foreground-prepare] ' + json.dumps(changes), flush=True)


def _stage_foreground_models():
    """Verify and mirror only the pinned RMBG group, sharing existing blobs."""
    from hf_space.download_models import source_inventory
    from hf_space.weight_cache import register_sources, _mirror
    spec = json.loads((ROOT/'configs/pipeline.json').read_text())['external_models']['rmbg2']
    operation = dict(spec, kind='hub', name='rmbg2', destination='external/rmbg2')
    sources = source_inventory(operation, MODEL_ROOT)
    total = sum(record['size'] for record in sources)
    _update_foreground(startup_progress={'label':'Preparing background-removal model',
        'completed':0, 'total':total, 'unit':'bytes'})
    copied = {}
    def progress(source, completed, size):
        copied[str(source)] = completed
        _update_foreground(startup_progress={'label':'Preparing background-removal model',
            'completed':sum(copied.values()), 'total':total, 'unit':'bytes'})
    def retry(source, attempt, attempts):
        _update_foreground(startup_progress={'label':f'Reconnecting to model storage (attempt {attempt} of {attempts})',
            'completed':sum(copied.values()), 'total':total, 'unit':'bytes'})
    register_sources(sources, retry=retry)
    destination = Path('/tmp/supratex-foreground-models-v1')
    _mirror(MODEL_ROOT/'external/rmbg2', destination, progress)
    return destination


def prepare_foreground_models():
    """CPU-only, independent of native compilation and the 18-model pipeline."""
    try:
        _update_foreground(stage='staging_local_weights', error=None, startup_progress={
            'label':'Checking background-removal model files', 'completed':0, 'total':1, 'unit':'steps'})
        model_path = _stage_foreground_models()
        _update_foreground(stage='loading_models', startup_progress={
            'label':'Loading background-removal model', 'completed':0, 'total':1, 'unit':'models'})
        from hf_space import zerogpu_runtime
        def notify(label, completed, total):
            _update_foreground(startup_progress={'label':label, 'completed':completed,
                                                'total':total, 'unit':'models'})
        summary = zerogpu_runtime.preload_foreground(model_path, notify)
        _update_foreground(stage='ready_released', error=None, resident_models=summary, startup_progress={
            'label':'Background-removal model ready', 'completed':1, 'total':1, 'unit':'models'})
    except BaseException:
        _update_foreground(stage='failed', error=traceback.format_exc()[-8000:], startup_progress={
            'label':'Background-removal model preparation failed. Please try again later.',
            'completed':0, 'total':1, 'unit':'steps'})
        traceback.print_exc()


def start_foreground_preparation():
    global _foreground_worker, _foreground_started
    if not zero_gpu_enabled():
        return start_preparation()
    with _lock:
        if not _foreground_started:
            _foreground_started = True
            _foreground_worker = threading.Thread(target=prepare_foreground_models, daemon=True,
                                                  name='foreground-model-prepare')
            _foreground_worker.start()
        return _foreground_worker


def update(**changes):
    global _snapshot
    with _persist_lock:
        with _state_lock:
            value=_snapshot|changes|{'epoch':_epoch,'updated_unix':time.time()}
            _snapshot=value
        STATE.parent.mkdir(parents=True,exist_ok=True)
        temporary=STATE.with_suffix('.tmp')
        temporary.write_text(json.dumps(value,indent=2))
        temporary.replace(STATE)
    print('[release-prepare] '+json.dumps(changes),flush=True)

def _stage_models():
    """Model metadata, verified transfers and bindings run in a killable child."""
    try:
        repository=os.environ.get('SUPRATEX_MODEL_REPO','supra3d/supra-tex')
        revision=os.environ.get('SUPRATEX_MODEL_REVISION',MODEL_REVISION)
        if (repository,revision)!=('supra3d/supra-tex',MODEL_REVISION):
            raise RuntimeError('This source release requires the validated, pinned model revision')
        from scripts import prepare_models as prep
        from hf_space.download_models import download
        plan=prep.build_plan(prep.parse_args(['--repo',repository,'--revision',revision,
            '--with-rmbg','--output-dir',str(MODEL_ROOT)]),resolve_remote=True)
        operations=sorted(plan['downloads'],key=lambda item:not item.get('gated',False))
        sources=[]
        for index,operation in enumerate(operations):
            if operation['name']=='project':
                operation=dict(operation,destination='project/'+revision)
            update(stage='downloading_foundations',project_revision=revision,
                   startup_progress={'label':'Checking model files: '+operation['name'],
                       'completed':index,'total':len(operations),'unit':'groups'})
            # Metadata resolves to the very same pinned release. Startup no
            # longer rehashes or copies large checkpoints through Bucket FUSE.
            sources.extend(download(operation,MODEL_ROOT,inventory_only=True))
        from hf_space.weight_cache import mirror_models,mirror_project,register_sources
        total=sum(record['size'] for record in sources)
        update(stage='staging_local_weights',startup_progress={
            'label':'Preparing local weight files','completed':0,'total':total,'unit':'bytes'})
        project_root=MODEL_ROOT/'project'/revision
        copied={}
        def file_progress(source,completed,size):
            copied[str(source)]=completed
            update(startup_progress={'label':'Preparing local weight files','completed':sum(copied.values()),
                'total':total,'unit':'bytes'})
        def retry_progress(source,attempt,attempts):
            update(startup_progress={'label':f'Reconnecting to model storage (attempt {attempt} of {attempts})',
                'completed':sum(copied.values()),'total':total,'unit':'bytes'})
        register_sources(sources,retry=retry_progress)
        foundations=mirror_models(MODEL_ROOT,progress=file_progress)
        project=mirror_project(project_root,revision,progress=file_progress)
        arguments=['--repo',str(project),'--output-dir',str(READY)]
        for flag,name in [('qwen-base','qwen-base'),('geometry','geometry'),('dino','dino'),
                          ('vae','vae'),('moge','moge2'),('rmbg','rmbg2')]:
            arguments+=['--'+flag,str(foundations/name)]
        arguments+=['--naf-weights',str(foundations/'naf/naf_release.pth')]
        config=prep.prepare(prep.build_plan(prep.parse_args(arguments)))
        prepared=json.loads(config.read_text())
        if not prepared['checkpoint'].endswith(CHECKPOINT):
            raise RuntimeError('The prepared pipeline did not select step-30000 completion weights')
        update(stage='bindings_ready',config=str(config),project_repository=repository,
               project_revision=revision,code_revision=CODE_REVISION,weight_origin='released',
               actual_training_steps=30000,random_weights_present=False,
               checkpoint=CHECKPOINT,model_loading='startup_pool_with_stage_offload',
               startup_progress={'label':'Initializing model runtime','completed':0,'total':1,'unit':'steps'})
        return config
    except BaseException:
        update(stage='failed',error=traceback.format_exc())
        traceback.print_exc()
        raise


def _read_child(path):
    try:
        value=json.loads(Path(path).read_text())
    except (FileNotFoundError,ValueError):
        return {}
    return value if value.get('epoch')==_epoch else {}


def _progress_key(value):
    # A periodic heartbeat or timestamp alone is not evidence of progress.
    progress=value.get('startup_progress',{})
    return (value.get('stage'),progress.get('label'),progress.get('completed'),progress.get('total'))


def _remaining(cap):
    if _shutdown.is_set():
        raise RuntimeError('Model startup was stopped')
    seconds=cap if _startup_deadline is None else min(cap,_startup_deadline-time.monotonic())
    if seconds<=0:
        raise TimeoutError('Model startup exceeded its overall time limit')
    return seconds


def _stage_attempt(attempt):
    global _staging_process
    from hf_space.startup_guard import process_options,supervise,stop_process
    seconds=_remaining(STAGING_SECONDS)
    folder=_runtime_directory/'preparation'/str(attempt)
    folder.mkdir(parents=True,exist_ok=True)
    status_file=folder/'status.json'
    status_file.unlink(missing_ok=True)
    env=native_build.runtime_environment()|{'SUPRATEX_STARTUP_EPOCH':_epoch,'PYTHONUNBUFFERED':'1'}
    process=None
    last=None
    def publish():
        nonlocal last
        value=_read_child(status_file)
        if value and value!=last:
            last=value
            fields={key:value[key] for key in value if key not in ('epoch','updated_unix','error')}
            if fields.get('stage')=='failed':
                fields.pop('stage')
            update(**fields)
    try:
        with (folder/'startup.log').open('w') as log:
            process=_staging_process=subprocess.Popen([
                sys.executable,'-m','hf_space.release_backend','--stage',str(status_file)],
                cwd=ROOT,env=env,stdin=subprocess.DEVNULL,stdout=log,stderr=subprocess.STDOUT,
                **process_options())
        code=supervise(process,total_seconds=seconds,idle_seconds=STAGING_IDLE_SECONDS,
            activity=lambda:_progress_key(_read_child(status_file)),tick=publish)
        publish()
        result=_read_child(status_file)
        if code or result.get('stage')!='bindings_ready':
            raise RuntimeError(result.get('error','Model file preparation failed'))
        return Path(result['config'])
    finally:
        # A retry is allowed only after this process is known to have exited.
        stop_process(process)
        _staging_process=None


def prepare_models():
    """The one boot controller owns native wait, retry, and runtime creation."""
    global _startup_deadline
    from hf_space.startup_guard import StartupCleanupError
    _startup_deadline=time.monotonic()+STARTUP_TOTAL_SECONDS
    try:
        update(stage='waiting_native_build',startup_progress={
            'label':'Preparing runtime','completed':0,'total':1,'unit':'steps'})
        native_build.start()
        deadline=time.monotonic()+_remaining(NATIVE_WAIT_SECONDS)
        while native_build.status().get('stage') not in ('compiled','failed'):
            if time.monotonic()>=deadline:
                raise TimeoutError('Native runtime preparation exceeded its startup deadline')
            time.sleep(.5)
        if native_build.status()['stage']!='compiled':
            raise RuntimeError('Native runtime preparation failed')
        for attempt in range(1,STARTUP_ATTEMPTS+1):
            try:
                _remaining(STARTUP_TOTAL_SECONDS)
                update(stage='staging_local_weights',attempt=attempt,attempts=STARTUP_ATTEMPTS,
                    error=None,startup_progress={
                        'label':'Checking model files','completed':0,'total':1,'unit':'steps'})
                config=_stage_attempt(attempt)
                update(stage='loading_models',startup_progress={
                    'label':'Initializing model runtime','completed':0,'total':1,'unit':'steps'})
                launch_runtime(config)
                update(stage='ready_released',error=None)
                _start_runtime_monitor()
                return
            except StartupCleanupError:
                raise
            except Exception:
                error=traceback.format_exc()[-8000:]
                stop_runtime()  # Must reap the previous CUDA process first.
                if attempt==STARTUP_ATTEMPTS:
                    raise
                update(stage='retrying',error=error,startup_progress={
                    'label':'Retrying model startup','completed':0,'total':1,'unit':'steps'})
                time.sleep(min(STARTUP_BACKOFF_SECONDS*attempt,_remaining(STARTUP_TOTAL_SECONDS)))
    except BaseException:
        update(stage='failed',error=traceback.format_exc()[-8000:],startup_progress={
            'label':'Model startup failed. Please try again later.','completed':0,'total':1,'unit':'steps'})
        traceback.print_exc()

def start_preparation():
    global _worker,_started
    with _lock:
        if not _started:
            _started=True
            _worker=threading.Thread(target=prepare_models,daemon=True,name='latest-release-prepare')
            _worker.start()
        return _worker


def stop_runtime():
    global _runtime_process
    if zero_gpu_enabled():
        runtime = sys.modules.get('hf_space.zerogpu_runtime')
        if runtime is not None:
            runtime.clear()
    process=_runtime_process
    if process is not None:
        from hf_space.startup_guard import stop_process
        stop_process(process)
        if _runtime_process is process:
            _runtime_process=None


def stop_preparation():
    _shutdown.set()
    if _staging_process is not None:
        from hf_space.startup_guard import stop_process
        stop_process(_staging_process)
    stop_runtime()


atexit.register(stop_preparation)


def launch_runtime(config):
    global _runtime_process
    if zero_gpu_enabled():
        from hf_space import zerogpu_runtime
        def notify(label, completed, total):
            update(startup_progress={'label': label, 'completed': completed,
                                     'total': total, 'unit': 'models'})
        summary = zerogpu_runtime.preload(config, notify)
        update(resident_models=summary, model_loading='cpu_pool_with_per_user_zerogpu',
               startup_progress={'label': 'Models ready',
                   'completed': summary['loaded_models'],
                   'total': summary['loaded_models'], 'unit': 'models'})
        return
    from hf_space.startup_guard import process_options
    if _runtime_process is not None:
        stop_runtime()
    seconds=_remaining(RUNTIME_SECONDS)
    _runtime_directory.mkdir(parents=True,exist_ok=True)
    status_file=_runtime_directory/'status.json'
    status_file.unlink(missing_ok=True)
    environment=native_build.runtime_environment()
    environment.update(OMP_NUM_THREADS='4',MKL_NUM_THREADS='4',
        SUPRATEX_RUNTIME_PROFILE='torch28',PYTHONUNBUFFERED='1',
        TOKENIZERS_PARALLELISM='false')
    with (_runtime_directory/'startup.log').open('w') as log:
        _runtime_process=subprocess.Popen([
            sys.executable,str(ROOT/'hf_space/runtime_worker.py'),'--config',str(config),
            '--status',str(_runtime_directory/'status.json'),'--epoch',_epoch],
            cwd=ROOT,env=environment,stdin=subprocess.PIPE,stdout=log,stderr=subprocess.STDOUT,
            text=True,**process_options())
    began=last_progress=time.monotonic()
    previous=None
    while time.monotonic()-began<seconds:
        if _runtime_process.poll() is not None:
            raise RuntimeError('Background model loading failed: '+
                (_runtime_directory/'startup.log').read_text(errors='replace')[-5000:])
        state=_read_child(status_file)
        if state:
            current=_progress_key(state)
            if current!=previous:
                previous=current
                last_progress=time.monotonic()
                update(**{key:state[key] for key in ('startup_progress','resident_models') if key in state})
            if state.get('stage')=='failed':
                raise RuntimeError(state.get('error','Background model loading failed'))
            if state.get('stage')=='ready':
                return
        if time.monotonic()-last_progress>=RUNTIME_IDLE_SECONDS:
            raise TimeoutError('Background model loading made no progress within its startup deadline')
        time.sleep(.5)
    raise TimeoutError('Background model loading exceeded its startup deadline')


def _start_runtime_monitor():
    global _runtime_monitor
    if zero_gpu_enabled():
        # There is deliberately no resident CUDA subprocess to monitor.
        return
    process=_runtime_process
    def monitor():
        last=None
        while _runtime_process is process:
            if process.poll() is not None:
                update(stage='failed',error='The model worker stopped. Restart the Space to reload models.')
                return
            state=_read_child(_runtime_directory/'status.json')
            fields={key:state[key] for key in ('startup_progress','resident_models') if key in state}
            if state.get('stage')=='failed':
                update(stage='failed',error=state.get('error','The model worker stopped.'))
                stop_runtime()
                return
            if fields and fields!=last:
                last=fields
                update(**fields)
            time.sleep(1)
    _runtime_monitor=threading.Thread(target=monitor,name='model-runtime-monitor',daemon=True)
    _runtime_monitor.start()

def ensure_prepared(generation_token=None):
    worker=start_preparation()
    while worker is not None and worker.is_alive():
        state=model_status()
        generation_progress.update(generation_token,status='preparing',
            label=state.get('startup_progress',{}).get('label','Preparing model files'))
        worker.join(timeout=1)
    state=model_status()
    if state.get('stage')!='ready_released':
        raise RuntimeError('Latest-release preparation failed: '+state.get('error','unknown failure')[-2000:])


def run_resident(job, seed, timeout_seconds, generation_token, *, operation='generate', foreground_prepared=False):
    zero_gpu = zero_gpu_enabled()
    # Each ZeroGPU lease owns an isolated fork. Only the dedicated runtime's
    # single stdin/result protocol needs a process-wide execution lock.
    with nullcontext() if zero_gpu else _generation_lock:
        if operation not in ('generate','remove_background','generate_geometry','generate_texture') or type(foreground_prepared) is not bool:
            raise ValueError('Invalid internal model operation')
        state = foreground_status() if zero_gpu and operation == 'remove_background' else model_status()
        if state.get('stage')!='ready_released':
            raise RuntimeError('Background model loading is incomplete')
        if zero_gpu:
            from hf_space import zerogpu_runtime
            if operation == 'remove_background':
                return zerogpu_runtime.foreground(str(job))
            if operation == 'generate_geometry':
                return zerogpu_runtime.geometry(str(job), int(seed), foreground_prepared)
            if operation == 'generate_texture':
                return zerogpu_runtime.texture(str(job), int(seed))
            return zerogpu_runtime.generate(str(job), int(seed), foreground_prepared)
        process=_runtime_process
        if process is None or process.poll() is not None:
            raise RuntimeError('The preloaded model worker is unavailable')
        request={'epoch':_epoch,'job':str(job),'seed':int(seed),
                 'operation':operation,'foreground_prepared':foreground_prepared}
        process.stdin.write(json.dumps(request)+'\n')
        process.stdin.flush()
        response_path=job/'resident-result.json'
        deadline=time.monotonic()+timeout_seconds
        last_stage=None
        while not response_path.exists():
            if process.poll() is not None:
                raise RuntimeError('The model worker stopped during '+operation.replace('_',' '))
            if time.monotonic()>deadline:
                stop_runtime()
                label='Background removal' if operation=='remove_background' else 'Generation'
                raise TimeoutError(label+' timed out; model worker stopped and requires a Space restart')
            receipt=job/'asset_work/run.json'
            if receipt.exists():
                stage=json.loads(receipt.read_text()).get('stage')
                if stage!=last_stage:
                    last_stage=stage
                    generation_progress.update(generation_token,label={
                        'frontend':'Generating geometry and reference materials',
                        'completion':'Completing surface materials and rendering o-voxel previews',
                    }.get(stage,'Starting generation…'))
            time.sleep(.5)
        response=json.loads(response_path.read_text())
        if response.get('epoch')!=_epoch:
            raise RuntimeError('Result belongs to another model runtime')
        if response.get('operation',operation)!=operation:
            raise RuntimeError('Result belongs to another model operation')
        if not response.get('ok'):
            label='Background removal failed: ' if operation=='remove_background' else 'Latest single-view generation failed: '
            raise RuntimeError(label+response.get('error','Unknown error')[-5000:])
        return response['resident_models']


def remove_background(image,timeout_seconds=120):
    """Run the queued cutout in the one startup-loaded, serial model worker."""
    if foreground_status().get('stage')!='ready_released':
        raise RuntimeError('Model files are not ready')
    job=JOBS_ROOT/uuid.uuid4().hex
    job.mkdir(parents=True)
    image.convert('RGBA').save(job/'input.png')
    try:
        run_resident(job,0,timeout_seconds,None,operation='remove_background')
        from PIL import Image
        with Image.open(job/'foreground.png') as result:
            if result.mode!='RGBA':
                raise RuntimeError('Background removal did not return an RGBA foreground')
            result.load()
            foreground=result.copy()
        metadata=json.loads((job/'foreground.json').read_text())
        if metadata.get('uncropped') is not True or metadata.get('source') not in ('rmbg2','input_alpha'):
            raise RuntimeError('Background removal returned invalid provenance')
        return foreground,metadata
    finally:
        # This transient worker exchange is not an asset/history directory.
        import shutil
        shutil.rmtree(job,ignore_errors=True)


def generate(image,seed,progress=None,timeout_seconds=1800,generation_token=None,*,foreground_prepared=False):
    if type(foreground_prepared) is not bool:
        raise ValueError('Prepared foreground must be an internal Boolean')
    state=model_status()
    if state.get('stage')!='ready_released':raise RuntimeError('Model files are not ready')
    identifier=uuid.uuid4().hex
    job=JOBS_ROOT/identifier
    job.mkdir(parents=True)
    image.save(job/'input.png')
    output=job/'asset.npz'
    started=time.monotonic()
    generation_progress.update(generation_token,status='running',job=identifier,
        generation_started_unix=time.time(),label='Starting generation…')
    if progress:progress(.05,'Starting generation…')
    resident_models=run_resident(job,seed,timeout_seconds,generation_token,foreground_prepared=foreground_prepared)
    receipt_path=job/'asset_work/run.json'
    receipt=json.loads(receipt_path.read_text()) if receipt_path.exists() else {}
    if receipt.get('status')!='complete' or not output.is_file():
        raise RuntimeError('Latest single-view generation failed: '+(job/'inference.log').read_text(errors='replace')[-5000:])
    import numpy as np
    with np.load(output,allow_pickle=False) as asset:
        required={'vertices','faces','coords','attrs','resolution','origin','voxel_size'}
        if not required<=set(asset.files) or not len(asset['coords']) or not len(asset['faces']):
            raise RuntimeError('Generated o-voxel geometry is incomplete')
        if asset['attrs'].shape[1]!=6 or not np.isfinite(asset['attrs']).all():
            raise RuntimeError('Generated o-voxel materials are invalid')
        voxel_count=len(asset['coords'])
    preview=json.loads((job/'preview.json').read_text())
    if len(list((job/'views').glob('*.webp')))!=preview['views']*len(preview['modes']):
        raise RuntimeError('The o-voxel preview render set is incomplete')
    if list(job.rglob('*.glb')) or preview.get('automatic_glb_export') is not False:
        raise RuntimeError('Generation must not export a mesh before selection')
    report={'status':'complete','asset_id':identifier,'code_revision':CODE_REVISION,
        'project_revision':MODEL_REVISION,'checkpoint':CHECKPOINT,'weight_origin':'released',
        'elapsed_seconds':time.monotonic()-started,'output_kind':'final_ovoxel','backview':False,
        'refine':False,'model_loading':('cpu_pool_with_per_user_zerogpu' if zero_gpu_enabled()
            else 'startup_pool_with_stage_offload'),'resident_models':resident_models,'run':receipt,
        'preview':preview,'voxel_count':voxel_count,'automatic_glb_export':False,
        'foreground_prepared':foreground_prepared}
    (job/'report.json').write_text(json.dumps(report,indent=2))
    if not (job/'timings.jsonl').exists():
        (job/'timings.jsonl').write_text(json.dumps({'name':'preview_generation','elapsed_seconds':report['elapsed_seconds']})+'\n')
    if progress:progress(1.,'Complete')
    return str(output),str(output),report


def _split_job(generation_token, label):
    if model_status().get('stage') != 'ready_released':
        raise RuntimeError('Model files are not ready')
    job = JOBS_ROOT / uuid.uuid4().hex
    job.mkdir(parents=True)
    generation_progress.update(generation_token, status='running', job=job.name,
        generation_started_unix=time.time(), label=label)
    return job


def _split_result(job, started, resident_models, *, geometry=False):
    import numpy as np
    result = json.loads((job / 'stage-result.json').read_text(encoding='utf-8'))
    expected = 'geometry' if geometry else 'final_ovoxel'
    if result.get('status') != 'complete' or result.get('output_kind') != expected:
        raise RuntimeError('The generation stage did not complete')
    output = job / ('geometry.npz' if geometry else 'asset.npz')
    with np.load(output, allow_pickle=False) as asset:
        required = {'vertices', 'faces'} if geometry else {'vertices','faces','coords','attrs','resolution','origin','voxel_size'}
        if not required <= set(asset.files) or not len(asset['vertices']) or not len(asset['faces']):
            raise RuntimeError('The generated geometry is incomplete')
        if not geometry and (not len(asset['coords']) or asset['attrs'].shape[1] != 6
                             or not np.isfinite(asset['attrs']).all()):
            raise RuntimeError('The generated material is incomplete')
        face_count = len(asset['faces'])
        voxel_count = None if geometry else len(asset['coords'])
    if geometry:
        from hf_space.split_pipeline import CHECKPOINT_DIRECTORY, validate_checkpoint
        validate_checkpoint(job / CHECKPOINT_DIRECTORY)
    preview = json.loads((job / 'preview.json').read_text(encoding='utf-8'))
    expected_frames = {f'{mode}-{index:02d}.webp' for mode in preview['modes']
                       for index in range(preview['views'])}
    frames = list((job / 'views').glob('*.webp'))
    if not expected_frames or {path.name for path in frames} != expected_frames or any(not p.stat().st_size for p in frames):
        raise RuntimeError('The preview render set is incomplete')
    if geometry and (preview['modes'] != ['clay'] or preview.get('default_mode') != 'clay'):
        raise RuntimeError('The geometry preview must contain clay views only')
    if list(job.rglob('*.glb')) or preview.get('automatic_glb_export') is not False:
        raise RuntimeError('Generation must not export a mesh before selection')
    report = {'status':'complete', 'asset_id':job.name, 'code_revision':CODE_REVISION,
        'project_revision':MODEL_REVISION, 'checkpoint':CHECKPOINT, 'weight_origin':'released',
        'elapsed_seconds':time.monotonic()-started, 'output_kind':expected,
        'backview':False, 'refine':False, 'automatic_glb_export':False,
        'model_loading':('cpu_pool_with_per_user_zerogpu' if zero_gpu_enabled()
                         else 'startup_pool_with_stage_offload'),
        'resident_models':resident_models, 'run':result, 'preview':preview,
        'face_count':face_count, 'geometry_reused':not geometry}
    if voxel_count is not None:
        report['voxel_count'] = voxel_count
        admission = output.with_suffix('.observation_admission.json')
        report['depth_pre_filter'] = {
            'code_revision': TEXTURE_FILTER_REVISION,
            'admission': json.loads(admission.read_text(encoding='utf-8')) if admission.is_file() else None,
        }
    (job / 'report.json').write_text(json.dumps(report, indent=2), encoding='utf-8')
    if not (job / 'timings.jsonl').exists():
        (job / 'timings.jsonl').write_text('', encoding='utf-8')
    return str(output), str(output), report


def generate_geometry(image, seed, generation_token=None, *, foreground_prepared=False, timeout_seconds=1800):
    if type(foreground_prepared) is not bool:
        raise ValueError('Prepared foreground must be an internal Boolean')
    started = time.monotonic()
    job = _split_job(generation_token, 'Generating geometry')
    image.save(job / 'input.png')
    summary = run_resident(job, int(seed), timeout_seconds, generation_token,
        operation='generate_geometry', foreground_prepared=foreground_prepared)
    return _split_result(job, started, summary, geometry=True)


def generate_texture(geometry_folder, seed, generation_token=None, *, timeout_seconds=1800):
    import shutil
    from hf_space.split_pipeline import CHECKPOINT_DIRECTORY, copy_checkpoint
    geometry_folder = Path(geometry_folder)
    started = time.monotonic()
    job = _split_job(generation_token, 'Generating materials')
    # Copy from the verified durable checkpoint before acquiring a GPU lease.
    copy_checkpoint(geometry_folder / CHECKPOINT_DIRECTORY, job / CHECKPOINT_DIRECTORY, seed=seed)
    shutil.copyfile(geometry_folder / 'input.png', job / 'input.png')
    summary = run_resident(job, int(seed), timeout_seconds, generation_token, operation='generate_texture')
    return _split_result(job, started, summary)


if __name__=='__main__':
    if len(sys.argv)!=3 or sys.argv[1]!='--stage':
        raise SystemExit('This module is an internal model preparation worker')
    STATE=Path(sys.argv[2])
    _stage_models()