"""Owned, serial ComfyUI test jobs. Never talks to the production image router."""
import json
import os
from pathlib import Path
import secrets
import shutil
import signal
import socket
import subprocess
import threading
import time
import urllib.request
import uuid
from profiles import QWEN_REPO, FLUX_REPO, image_recipe

PYTHON=Path('/opt/deck-image-python/bin/python')
COMFY=Path('/opt/deck-comfy')
GIB=1024**3

def cgroup_headroom():
    """Exclude reclaimable inactive file cache, not live anonymous allocations."""
    try:
        raw=Path('/sys/fs/cgroup/memory.max').read_text().strip()
        if raw=='max':return None
        used=int(Path('/sys/fs/cgroup/memory.current').read_text())
        stat=dict(line.split() for line in Path('/sys/fs/cgroup/memory.stat').read_text().splitlines())
        return int(raw)-max(0,used-int(stat.get('inactive_file',0)))
    except FileNotFoundError:return None

def probe():
    rows=subprocess.run(['nvidia-smi','--query-gpu=uuid,name,memory.total,memory.free','--format=csv,noheader,nounits'],capture_output=True,text=True,check=True,timeout=5).stdout
    apps=subprocess.run(['nvidia-smi','--query-compute-apps=gpu_uuid,pid','--format=csv,noheader,nounits'],capture_output=True,text=True,check=True,timeout=5).stdout
    counts={}
    for line in apps.splitlines():
        gpu,_=line.split(',',1);counts[gpu.strip()]=counts.get(gpu.strip(),0)+1
    return [dict(uuid=u.strip(),name=n.strip(),total_mib=float(t),free_mib=float(f),processes=counts.get(u.strip(),0)) for u,n,t,f in (line.split(',') for line in rows.splitlines())]

def select_gpus(devices,model_size,encoder_size,vae_size,offload=False):
    needed=(min(model_size,8*GIB)+vae_size+3*GIB if offload else model_size+vae_size+3*GIB)/1024**2
    images=[g for g in devices if 'RTX 5080' in g['name'] and g['processes']==0 and g['free_mib']>needed]
    encoders=[g for g in devices if 'RTX 3060' in g['name'] and g['processes']==0 and g['free_mib']>(encoder_size+2*GIB)/1024**2]
    if not encoders:raise ValueError('Textencoder benötigt eine freie RTX 3060 mit ausreichend VRAM. Kein automatischer CPU-Fallback.')
    if not images:raise ValueError('Bildmodell benötigt eine freie RTX 5080 mit ausreichend VRAM. Deck stoppt keine anderen Modelle.')
    return max(images,key=lambda g:g['free_mib']),encoders[0]

def workflow(prompt,params,seed,family='qwen'):
    if family=='flux':
        return {
          '1':{'class_type':'UNETLoader','inputs':{'unet_name':'model.safetensors','weight_dtype':'default'}},
          '2':{'class_type':'DeckTextEncoderLoader','inputs':{'clip_name':'encoder.safetensors','family':'flux2'}},
          '3':{'class_type':'VAELoader','inputs':{'vae_name':'vae.safetensors'}},
          '4':{'class_type':'CLIPTextEncode','inputs':{'clip':['2',0],'text':prompt}},
          '5':{'class_type':'EmptyFlux2LatentImage','inputs':{'width':params['width'],'height':params['height'],'batch_size':1}},
          '6':{'class_type':'SamplerCustomAdvanced','inputs':{'noise':['9',0],'guider':['10',0],'sampler':['11',0],'sigmas':['12',0],'latent_image':['5',0]}},
          '7':{'class_type':'VAEDecode','inputs':{'samples':['6',0],'vae':['3',0]}},
          '8':{'class_type':'SaveImage','inputs':{'filename_prefix':'result','images':['7',0]}},
          '9':{'class_type':'RandomNoise','inputs':{'noise_seed':seed}},
          '10':{'class_type':'CFGGuider','inputs':{'model':['1',0],'positive':['4',0],'negative':['13',0],'cfg':params['guidance']}},
          '11':{'class_type':'KSamplerSelect','inputs':{'sampler_name':'euler'}},
          '12':{'class_type':'Flux2Scheduler','inputs':{'steps':params['steps'],'width':params['width'],'height':params['height']}},
          '13':{'class_type':'ConditioningZeroOut','inputs':{'conditioning':['4',0]}}
        }

    return {
      '1':{'class_type':'UnetLoaderGGUF','inputs':{'unet_name':'model.gguf'}},
      '2':{'class_type':'DeckTextEncoderLoader','inputs':{'clip_name':'encoder.safetensors'}},
      '3':{'class_type':'VAELoader','inputs':{'vae_name':'vae.safetensors'}},
      '4':{'class_type':'TextEncodeQwenImage21','inputs':{'clip':['2',0],'prompt':prompt,'negative_prompt':'','resolution':max(params['width'],params['height'])}},
      '5':{'class_type':'EmptyLatentImage','inputs':{'width':params['width'],'height':params['height'],'batch_size':1}},
      '6':{'class_type':'KSampler','inputs':{'model':['1',0],'positive':['4',0],'negative':['4',1],'latent_image':['5',0],'seed':seed,'steps':params['steps'],'cfg':params['guidance'],'sampler_name':'euler','scheduler':'simple','denoise':1}},
      '7':{'class_type':'VAEDecode','inputs':{'samples':['6',0],'vae':['3',0]}},
      '8':{'class_type':'SaveImage','inputs':{'filename_prefix':'result','images':['7',0]}}
    }

class ImageTests:
 def __init__(self,root,profiles):
    self.root=Path(root);self.profiles=profiles;self.runtime=None;self.lock=threading.RLock();self.process=None;self.cancel=threading.Event();self.job=None;self.acquire=lambda wait=False:lambda:None
    if (self.root/'status.json').exists():
        self.job=json.loads((self.root/'status.json').read_text())
        if self.job.get('state')=='running':self.job.update(state='interrupted',phase='Deck wurde neu gestartet; Test erneut starten.')
 def _save(self):
    self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
    p=self.root/'status.tmp';p.write_text(json.dumps(self.job));p.replace(self.root/'status.json')
 def _phase(self,text):
    with self.lock:
        self.job['phase']=text;self._save()
 def status(self):
    with self.lock:
        python,comfy=self.runtime.paths() if self.runtime else (PYTHON,COMFY)
        return dict(job=dict(self.job) if self.job else None,runtime_installed=python.is_file() and (comfy/'main.py').is_file())
 def stop(self):
    self.cancel.set()
    with self.lock:process=self.process
    if process and process.poll() is None:
        try:
            os.killpg(process.pid,signal.SIGTERM)
            try:process.wait(timeout=3)
            except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait()
        except ProcessLookupError:pass
    return {'cancellation_requested':True}
 def model_path(self,item):
    return (self.profiles.catalog.root/item['id']/('model'+Path(item['file']).suffix)).resolve()
 def start(self,profile_id,prompt,wait=False,reserved=False):
    if not isinstance(prompt,str) or not 1<=len(prompt.strip())<=4000:raise ValueError('Bitte einen Prompt mit 1–4000 Zeichen eingeben.')
    release=(lambda:None) if reserved else self.acquire(wait)
    try:return self._start(profile_id,prompt,release)
    except Exception:release();raise
 def _start(self,profile_id,prompt,release):
    if not isinstance(prompt,str) or not 1<=len(prompt.strip())<=4000:raise ValueError('Bitte einen Prompt mit 1–4000 Zeichen eingeben.')
    with self.lock:
        if self.job and self.job['state']=='running':raise ValueError('Ein Bildtest läuft bereits.')
        if self.runtime and (self.runtime.status()['job'] or {}).get('state')=='running':raise ValueError('Die Bildlaufzeit wird gerade installiert.')
        if not self.status()['runtime_installed']:raise ValueError('Die eigene ComfyUI-Bildlaufzeit ist noch nicht installiert.')
        profile=next((p for p in self.profiles.status()['profiles'] if p['id']==profile_id),None)
        if not profile or profile['kind']!='image' or not profile['model']:raise ValueError('Bildprofil nicht verfügbar.')
        model=profile['model']
        if not image_recipe(model):raise ValueError('Für dieses Bildmodell fehlt ein Komponentenrezept.')
        params=profile['parameters']
        if params['width']>1024 or params['height']>1024:raise ValueError('Der isolierte Test unterstützt maximal 1024 × 1024 Pixel.')
        encoder=self.profiles._component(model,'text_encoder',profile.get('components',{}).get('text_encoder'))
        vae=self.profiles._component(model,'vae',profile.get('components',{}).get('vae'))
        mem={line.split(':')[0]:int(line.split()[1])*1024 for line in Path('/proc/meminfo').read_text().splitlines() if line.startswith(('MemAvailable:','MemTotal:'))}
        required_ram=max(20*GIB,encoder['size']*2+4*GIB,model['size']+8*GIB if model['repo']==FLUX_REPO else 0)
        headroom=cgroup_headroom()
        if headroom is not None and headroom<required_ram:raise ValueError('Deck-RAM-Limit reicht für Textencoder und Arbeitsdaten nicht aus.')
        if mem.get('MemAvailable',0)<required_ram+4*GIB:raise ValueError('Aktuell zu wenig freier System-RAM; produktive Dienste bleiben unverändert.')
        gpu,encoder_gpu=select_gpus(probe(),model['size'],encoder['size'],vae['size'],offload=model['repo']==FLUX_REPO);self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
        if shutil.disk_usage(self.root).free<10*GIB:raise ValueError('Weniger als 10 GiB freier Plattenspeicher.')
        job_id=uuid.uuid4().hex;seed=params['seed'] if params['seed']>=0 else secrets.randbelow(2147483648)
        self.cancel.clear();self.job=dict(id=job_id,state='running',phase='Bildlaufzeit startet',profile_id=profile_id,profile_name=profile['name'],started_at=time.time(),gpu=gpu['name'],encoder_gpu=encoder_gpu['name'],seed=seed,parameters=params)
        self._save();threading.Thread(target=self._run,args=(job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release),daemon=True).start()
        return dict(self.job)
 def _run(self,job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release=lambda:None):
    directory=self.root/job_id;process=None
    try:
        directory.mkdir(mode=0o700)
        for role,item,filename in [('unet',model,'model.safetensors' if model['repo']==FLUX_REPO else 'model.gguf'),('clip',encoder,'encoder.safetensors'),('vae',vae,'vae.safetensors')]:
            dest=directory/'models'/role;dest.mkdir(parents=True);(dest/filename).symlink_to(self.model_path(item))
        for folder in ('output','temp','user','input'):(directory/folder).mkdir()
        custom=directory/'custom_nodes'/'deck_encoder';custom.mkdir(parents=True)
        shutil.copyfile(Path(__file__).parent/'image_encoder_node.py',custom/'__init__.py')
        config={'deck_nodes':{'base_path':str(directory),'custom_nodes':'custom_nodes'},'deck':{'base_path':str(directory/'models'),'unet':'unet','clip':'clip','vae':'vae'}}
        (directory/'paths.json').write_text(json.dumps(config)) # JSON is valid YAML.
        with socket.socket() as sock:sock.bind(('127.0.0.1',0));port=sock.getsockname()[1]
        env=dict(os.environ,CUDA_VISIBLE_DEVICES=gpu['uuid']+','+encoder_gpu['uuid'],OMP_NUM_THREADS='2',MKL_NUM_THREADS='2',HOME=str(directory),HF_HUB_OFFLINE='1',TRANSFORMERS_OFFLINE='1',PYTHONDONTWRITEBYTECODE='1')
        args=[str(self.runtime.paths()[0] if self.runtime else PYTHON),str((self.runtime.paths()[1] if self.runtime else COMFY)/'main.py'),'--listen','127.0.0.1','--port',str(port),'--disable-auto-launch','--disable-metadata','--lowvram','--reserve-vram','1.5','--extra-model-paths-config',str(directory/'paths.json'),'--output-directory',str(directory/'output'),'--temp-directory',str(directory/'temp'),'--user-directory',str(directory/'user'),'--input-directory',str(directory/'input')]
        with self.lock:
            if self.cancel.is_set():raise InterruptedError()
            process=subprocess.Popen(args,cwd=self.runtime.paths()[1] if self.runtime else COMFY,env=env,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=True);self.process=process
        def request(path,data=None):
            req=urllib.request.Request(f'http://127.0.0.1:{port}'+path,data=json.dumps(data).encode() if data is not None else None,headers={'Content-Type':'application/json'})
            with urllib.request.urlopen(req,timeout=5) as r:
                raw=r.read(4*1024*1024+1)
                if len(raw)>4*1024*1024:raise ValueError('Bildlaufzeit-Antwort zu groß.')
                return json.loads(raw)
        def guard():
            if self.cancel.is_set():raise InterruptedError()
            if process.poll() is not None:raise ValueError('Bildlaufzeit wurde beendet (Speicherlimit oder Startfehler).')
            current={g['uuid']:g for g in probe()}
            for reserved in (gpu,encoder_gpu):
                row=current.get(reserved['uuid'])
                if not row or row['processes']>1:raise ValueError('Reservierte GPU wird inzwischen von einem weiteren Prozess verwendet. Deck-Test beendet.')
        deadline=time.monotonic()+180
        while True:
            guard()
            try:request('/system_stats');break
            except (OSError,ValueError):
                if time.monotonic()>deadline:raise ValueError('Bildlaufzeit wurde nicht rechtzeitig bereit.')
                time.sleep(1)
        nodes=request('/object_info')
        graph=workflow(prompt,params,seed,'flux' if model['repo']==FLUX_REPO else 'qwen')
        required={node['class_type'] for node in graph.values()}
        if not required.issubset(nodes):raise ValueError('Der installierten Bildlaufzeit fehlen erforderliche Nodes für dieses Bildmodell.')
        self._phase('Auftrag wird verarbeitet · Textencoder auf RTX 3060 · Bildberechnung auf RTX 5080')
        response=request('/prompt',{'prompt':graph,'client_id':'deck-'+job_id});prompt_id=response['prompt_id']
        deadline=time.monotonic()+1800
        while time.monotonic()<deadline:
            guard();history=request('/history/'+prompt_id)
            if prompt_id in history:
                result=history[prompt_id]
                if result.get('status',{}).get('status_str')=='error':
                    errors=[m[1] for m in result.get('status',{}).get('messages',[]) if m[0]=='execution_error']
                    error=errors[0] if errors else {}
                    # Do not include prompts, tensor dumps or upstream exception messages.
                    raise ValueError('Generierung fehlgeschlagen: '+str(error.get('node_type','Worker'))+' / '+str(error.get('exception_type','Fehler')))
                images=[i for value in result.get('outputs',{}).values() for i in value.get('images',[])]
                if images:
                    image=images[0];path=(directory/'output'/image.get('subfolder','')/image['filename']).resolve()
                    if not path.is_relative_to((directory/'output').resolve()) or path.stat().st_size>32*1024**2:raise ValueError('Ungültiges Ergebnisbild.')
                    with path.open('rb') as f:
                        if f.read(8)!=b'\x89PNG\r\n\x1a\n':raise ValueError('Ergebnis ist kein PNG.')
                    shutil.copyfile(path,directory/'result.png');break
            time.sleep(2)
        else:raise ValueError('Zeitlimit der Bildgenerierung erreicht.')
        final_state='complete';phase='Bild fertig · Modell wurde entladen'
    except InterruptedError:final_state='cancelled';phase='Bildtest abgebrochen'
    except Exception as exc:final_state='failed';phase=str(exc) if isinstance(exc,ValueError) else 'Bildlaufzeit nicht erreichbar oder nicht bereit. Komponenten und Installation prüfen.'
    finally:
        try:
            if process and process.poll() is None:
                try:os.killpg(process.pid,signal.SIGTERM);process.wait(timeout=10)
                except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait()
                except ProcessLookupError:pass
            with self.lock:
                self.process=None;self.job.update(state=final_state,phase=phase,finished_at=time.time());self._save()
        finally:release()
 def image(self,job_id):
    with self.lock:
        if not self.job or self.job['id']!=job_id or self.job['state']!='complete':raise ValueError('Ergebnisbild nicht verfügbar.')
        return (self.root/job_id/'result.png').read_bytes()