"""Owned, serial ComfyUI test jobs. Never talks to the production image router."""
import json
import base64
import binascii
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, image_capabilities

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',references=()):
    if family not in ('qwen','flux'):raise ValueError('Für diese Modellfamilie fehlt eine Bildlaufzeit-Anbindung.')
    if len(references)>4:raise ValueError('Maximal vier Referenzbilder pro Auftrag.')
    if family=='flux':
        graph={
          '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]}}
        }

        # FLUX.2 editing uses reference latents on both conditioning branches,
        # with fresh output noise; never silently discard uploaded references.
        positive=['4',0];negative=['13',0]
        for index,name in enumerate(references):
            load,scale,encode,pos,neg=map(str,range(20+index*5,25+index*5))
            graph[load]={'class_type':'LoadImage','inputs':{'image':name}}
            graph[scale]={'class_type':'ImageScaleToTotalPixels','inputs':{'image':[load,0],'upscale_method':'lanczos','megapixels':1.0,'resolution_steps':1}}
            graph[encode]={'class_type':'VAEEncode','inputs':{'pixels':[scale,0],'vae':['3',0]}}
            graph[pos]={'class_type':'ReferenceLatent','inputs':{'conditioning':positive,'latent':[encode,0]}}
            graph[neg]={'class_type':'ReferenceLatent','inputs':{'conditioning':negative,'latent':[encode,0]}}
            positive=[pos,0];negative=[neg,0]
        graph['10']['inputs'].update(positive=positive,negative=negative)
        return graph

    graph={
      '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]}}
    }
    if references:
        graph['4']['inputs']['vae']=['3',0]
        for index,name in enumerate(references,1):
            graph[str(20+index)]={'class_type':'LoadImage','inputs':{'image':name}}
            graph['4']['inputs'][f'images.image_{index}']=[str(20+index),0]
        graph['6']['inputs']['latent_image']=['4',2]
    return graph

def decode_references(values,limit):
    if not isinstance(values,list) or len(values)>limit:raise ValueError(f'Dieses Bildprofil erlaubt höchstens {limit} Referenzbilder.')
    result=[]
    for value in values:
        if not isinstance(value,str) or not value.startswith(('data:image/png;base64,','data:image/jpeg;base64,','data:image/webp;base64,')) or len(value)>14*1024*1024:raise ValueError('Referenzbild muss PNG, JPEG oder WebP bis 10 MiB sein.')
        try:raw=base64.b64decode(value.split(',',1)[1],validate=True)
        except (ValueError,binascii.Error):raise ValueError('Ungültige Bildkodierung.') from None
        if not 0<len(raw)<=10*1024*1024:raise ValueError('Referenzbild muss 1 bis 10 MiB groß sein.')
        kind=value.split(';',1)[0].split('/')[-1]
        signatures={'png':b'\x89PNG\r\n\x1a\n','jpeg':b'\xff\xd8\xff','webp':b'RIFF'}
        if not raw.startswith(signatures[kind]) or kind=='webp' and raw[8:12]!=b'WEBP':raise ValueError('Dateiinhalt entspricht nicht dem Bildformat.')
        result.append((raw,'jpg' if kind=='jpeg' else kind))
    return result

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,reference_images=None,skip_enhancer=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,reference_images or [],skip_enhancer)
    except Exception:release();raise
 def _start(self,profile_id,prompt,release,reference_images,skip_enhancer=False):
    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.')
        references=decode_references(reference_images,image_capabilities(model)['reference_images'])
        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,reference_count=len(references))
        self._save();threading.Thread(target=self._run,args=(job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release,references,profile,skip_enhancer),daemon=True).start()
        return dict(self.job)
 def _run(self,job_id,prompt,params,seed,gpu,encoder_gpu,model,encoder,vae,release=lambda:None,references=(),profile=None,skip_enhancer=False):
    directory=self.root/job_id;process=None
    try:
        if not skip_enhancer and profile and getattr(self,'enhancer',None) and (profile.get('prompt_enhancer') or {}).get('i2i' if references else 't2i'):
            self._phase('Prompt-Aufwerter lädt; Bildmodell wartet')
            result=self.enhancer.rewrite(profile,prompt,references,self.cancel)
            prompt=result['prompt']
            with self.lock:self.job['prompt_enhanced']=True;self.job['prompt_enhancer']=result['source'];self._save()
            if self.cancel.is_set():raise InterruptedError()
            self._phase('Prompt-Aufwerter entladen; Bildlaufzeit startet')
        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()
        names=[]
        for index,(raw,extension) in enumerate(references,1):
            name=f'reference-{index}.{extension}';(directory/'input'/name).write_bytes(raw);names.append(name)
        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',names)
        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()
            for name in (f'reference-{index}.{extension}' for index,(_,extension) in enumerate(references,1)):
                (directory/'input'/name).unlink(missing_ok=True)
        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()