#!/usr/bin/env python3
"""Install ComfyUI and dependencies on BC-250."""
import paramiko
import time

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect('192.168.178.150', username='fabian', key_filename=r'C:\Users\fabia\.ssh\id_ed25519')

def run(cmd, timeout=300, desc=""):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    print(f"$ {cmd}")
    _, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode()
    err = stderr.read().decode()
    rc = stdout.channel.recv_exit_status()
    if out.strip():
        lines = out.strip().split('\n')
        if len(lines) > 60:
            print(f"  ... ({len(lines)} lines, showing last 60)")
            print('\n'.join(lines[-60:]))
        else:
            print(out.strip())
    if err.strip():
        lines = err.strip().split('\n')
        show = lines[-20:] if len(lines) > 20 else lines
        print(f"STDERR: {chr(10).join(show)}")
    print(f"  Exit code: {rc}")
    return rc, out, err

# Clone ComfyUI
run("git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git ~/ComfyUI 2>&1 | tail -10",
    desc="Clone ComfyUI")

# Install ComfyUI requirements in venv
run("bash -c 'source ~/comfyui-env/bin/activate && cd ~/ComfyUI && pip install -r requirements.txt 2>&1 | tail -20'",
    desc="Install ComfyUI requirements",
    timeout=300)

# Install diffusers from source (needed for ZImagePipeline)
run("bash -c 'source ~/comfyui-env/bin/activate && pip install git+https://github.com/huggingface/diffusers 2>&1 | tail -10'",
    desc="Install diffusers from source (for ZImagePipeline)",
    timeout=300)

# Install additional deps that ComfyUI/Z-Image might need
run("bash -c 'source ~/comfyui-env/bin/activate && pip install transformers accelerate safetensors sentencepiece huggingface_hub aiohttp einops torchvision 2>&1 | tail -15'",
    desc="Install transformers, accelerate, etc.",
    timeout=300)

# Install ComfyUI-Manager (custom node manager)
run("git clone --depth 1 https://github.com/Comfy-Org/ComfyUI-Manager.git ~/ComfyUI/custom_nodes/ComfyUI-Manager 2>&1 | tail -5",
    desc="Install ComfyUI-Manager")

# Install ComfyUI-GGUF (needed for GGUF checkpoint format)
run("git clone --depth 1 https://github.com/city96/ComfyUI-GGUF.git ~/ComfyUI/custom_nodes/ComfyUI-GGUF 2>&1 | tail -5",
    desc="Install ComfyUI-GGUF nodes")

# Install GGUF dependencies
run("bash -c 'source ~/comfyui-env/bin/activate && pip install gguf 2>&1 | tail -5'",
    desc="Install gguf Python package")

# Install Z-Image Power Nodes
run("git clone --depth 1 https://github.com/martin-rizzo/ComfyUI-ZImagePowerNodes.git ~/ComfyUI/custom_nodes/ComfyUI-ZImagePowerNodes 2>&1 | tail -5",
    desc="Install Z-Image Power Nodes")

# Verify ComfyUI structure
run("ls -la ~/ComfyUI/main.py ~/ComfyUI/custom_nodes/ 2>&1",
    desc="Verify ComfyUI structure")

run("ls ~/ComfyUI/custom_nodes/",
    desc="Custom nodes installed")

# Create model directories
run("mkdir -p ~/ComfyUI/models/diffusion_models ~/ComfyUI/models/text_encoders ~/ComfyUI/models/vae ~/ComfyUI/models/checkpoints",
    desc="Create model directories")

# Quick test: can ComfyUI import?
run("""bash -c 'source ~/comfyui-env/bin/activate && \
    HSA_OVERRIDE_GFX_VERSION=10.1.0 \
    HIP_VISIBLE_DEVICES=0 \
    HSA_ENABLE_SDMA=0 \
    cd ~/ComfyUI && python3 -c "
import torch
print(f\\"torch {torch.__version__} hip={torch.version.hip} cuda={torch.cuda.is_available()}\\")
import comfy
print(\\"ComfyUI import OK\\")
" 2>&1'""",
    desc="Test ComfyUI import",
    timeout=60)

ssh.close()
print("\nDone — ComfyUI installed.")