"""Explicit optional SwarmUI service. Run on Athena as root after image build."""
import json,os,secrets,socket,subprocess,sys
from pathlib import Path
base=Path('/opt/athena-swarm-ui');deck=Path('/opt/athena-deck-dev/runtime/state');name='athena-swarm-ui'
def run(*args):return subprocess.check_output(args,text=True).strip()
existing=subprocess.run(['docker','inspect',name],capture_output=True,text=True)
if existing.returncode==0:
 item=json.loads(existing.stdout)[0];labels=item['Config'].get('Labels',{})
 if '--replace' not in sys.argv or labels.get('io.athena-deck.application')!='swarm-ui':raise SystemExit('Container exists. --replace only accepts the managed swarm-ui service.')
 run('docker','rm','-f',name)
for port in (8125,8126):
 with socket.socket() as sock:
  sock.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1);sock.bind(('127.0.0.1',port))
for folder in ('Data','Output','dlbackend'):
 p=base/folder;p.mkdir(parents=True,exist_ok=True);os.chown(p,65534,65534)
video=deck/'video';video.mkdir(exist_ok=True);token=video/'comfy-client-token'
if not token.exists():token.write_text(secrets.token_urlsafe(48)+'\n')
os.chmod(token,0o600);os.chown(token,65534,65534)
# Reuse Deck's ComfyUI; install only the official frontend nodes and OpenCV dependency.
nodes=video/'swarm-comfy-nodes';nodes.mkdir(exist_ok=True)
temporary=run('docker','create','athena-swarm-ui:de7b834')
try:run('docker','cp',temporary+':/swarm/src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/.',str(nodes))
finally:run('docker','rm',temporary)
subprocess.check_call(['docker','exec','athena-deck-dev','/opt/deck-image-python/bin/python','-m','pip','install','--upgrade','--no-cache-dir','--no-deps','--target','/var/lib/deck/video/swarm-python','opencv-python-headless==4.12.0.88','imageio-ffmpeg==0.6.0'])
for p in [nodes,*nodes.rglob('*')]:os.chown(p,65534,65534)
work=video/'comfy-work/models'
for folder in ('diffusion_models','text_encoders','vae','latent_upscale_models','loras','Stable-Diffusion','Lora','VAE','Embeddings','controlnet','model_patches','clip','clip_vision','upscale_models','tensorrt','unet'):
 p=work/folder;p.mkdir(parents=True,exist_ok=True);os.chown(p,65534,65534)
# Only link the maintained LTX recipe; no model copies or downloads.
entries=[]
for metadata in (deck/'models').glob('*/entry.json'):
 item=json.loads(metadata.read_text());item['id']=metadata.parent.name;entries.append(item)
main=next((x for x in entries if x.get('repo')=='Lightricks/LTX-2.5' and x.get('file')=='diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors'),None)
if main:
 for item in entries:
  if item.get('repo')!=main['repo'] or item.get('revision')!=main['revision']:continue
  filename=Path(item['file']);folder=filename.parts[0]
  if folder not in ('diffusion_models','text_encoders','vae','latent_upscale_models'):continue
  source=deck/'models'/item['id']/('model'+filename.suffix)
  if not source.is_file():continue
  link=work/folder/filename.name
  if not link.exists() and not link.is_symlink():link.symlink_to('/var/lib/deck/models/'+item['id']+'/model'+filename.suffix)
settings=base/'Data/Settings.fds'
if not settings.exists():settings.write_text('IsInstalled: true\nLaunchMode: none\nNetwork:\n    Host: 127.0.0.1\n    Port: 8125\n    PortCanChange: false\nPaths:\n    ModelRoot: Models\n    SDModelFolder: diffusion_models\n    SDVAEFolder: vae\n    SDLoraFolder: loras\n    SDClipFolder: text_encoders\n')
backends=base/'Data/Backends.fds'
if not backends.exists():backends.write_text('0:\n    type: comfyui_api\n    title: Athena Deck ComfyUI\n    enabled: true\n    settings:\n        Address: http://127.0.0.1:8126\n        AllowIdle: true\n        OverQueue: 0\n')
for p in (settings,backends):os.chown(p,65534,65534)
args=['docker','run','-d','--name',name,'--network','host','--restart','unless-stopped','--cap-drop','ALL','--security-opt','no-new-privileges:true','--memory','2g','--cpus','2','--label','io.athena-deck.managed=true','--label','io.athena-deck.role=application','--label','io.athena-deck.application=swarm-ui']
for folder in ('Data','Output','dlbackend'):args+=['-v',f'{base/folder}:/swarm/{folder}:rw']
args+=['-v',f'{work}:/swarm/Models:ro','-v',f'{deck/"models"}:/var/lib/deck/models:ro','-v',f'{token}:/run/secrets/comfy-token:ro','athena-swarm-ui:de7b834']
run(*args);print('SwarmUI installed: http://127.0.0.1:8125 · API By URL · no GPU runtime')