"""BC-250 Full GPU Fix: Verify ROCm, diagnose VRAM, start ComfyUI on GPU, generate image."""
import paramiko
import time
import json
import sys

SSH_HOST = '192.168.178.150'
SSH_USER = 'fabian'
SSH_KEY = r'C:\Users\fabia\.ssh\id_ed25519'

def ssh_connect():
    k = paramiko.Ed25519Key.from_private_key_file(SSH_KEY)
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(SSH_HOST, username=SSH_USER, pkey=k, timeout=15)
    return c

def run(c, cmd, timeout=30):
    """Run command via fish shell, return stdout."""
    wrapped = f'bash -c {repr(cmd)}'
    _, o, e = c.exec_command(wrapped, timeout=timeout)
    return o.read().decode(errors='replace').strip()

def run_full(c, cmd, timeout=30):
    """Run command, return (stdout, stderr)."""
    wrapped = f'bash -c {repr(cmd)}'
    _, o, e = c.exec_command(wrapped, timeout=timeout)
    return o.read().decode(errors='replace').strip(), e.read().decode(errors='replace').strip()

# ============================================================
# PHASE 1: Kill any remnants
# ============================================================
print("="*60)
print("PHASE 1: Clean slate")
print("="*60)
c = ssh_connect()
run(c, 'pkill -f "python.*main.py" 2>/dev/null; pkill -f comfyui 2>/dev/null')
time.sleep(2)
leftover = run(c, 'pgrep -af "python.*main.py" 2>/dev/null')
if leftover:
    print(f"WARNING: Still running: {leftover}")
    run(c, 'pkill -9 -f "python.*main.py" 2>/dev/null')
    time.sleep(1)
print("ComfyUI killed. Clean slate.")

# ============================================================
# PHASE 2: Verify ROCm + PyTorch GPU
# ============================================================
print("\n" + "="*60)
print("PHASE 2: Verify ROCm + PyTorch GPU access")
print("="*60)

# Set GPU env vars for ALL subsequent commands
GPU_ENV = (
    'export HSA_OVERRIDE_GFX_VERSION=10.1.0; '
    'export HIP_VISIBLE_DEVICES=0; '
    'export HSA_ENABLE_SDMA=0; '
    'export HSA_TOOLS_LIB=""; '
    'export HSA_TOOLS_REPORT_LOAD_FAILURE=0; '
    'export PYTORCH_HIP_ALLOC_CONF=expandable_segments:False; '
)

# Check rocminfo
out = run(c, f'{GPU_ENV} rocminfo 2>&1 | grep -E "Name:|Marketing Name:|gfx" | head -10')
print(f"ROCm devices:\n{out}")

# Check PyTorch GPU
gpu_test = f'''{GPU_ENV} cd ~/ComfyUI && source ~/comfyui-env/bin/activate.fish 2>/dev/null; . ~/comfyui-env/bin/activate 2>/dev/null; python3 -c "
import torch
print(f'PyTorch: {{torch.__version__}}')
print(f'CUDA available: {{torch.cuda.is_available()}}')
print(f'Device count: {{torch.cuda.device_count()}}')
if torch.cuda.is_available():
    print(f'Device name: {{torch.cuda.get_device_name(0)}}')
    free, total = torch.cuda.mem_get_info(0)
    print(f'VRAM: {{free//1024//1024}}MB free / {{total//1024//1024}}MB total')
    # Quick GPU compute test
    x = torch.randn(1024, 1024, device='cuda', dtype=torch.float16)
    y = torch.mm(x, x)
    print(f'GPU compute test: OK (result sum={{y.sum().item():.1f}})')
    del x, y
    torch.cuda.empty_cache()
else:
    print('ERROR: GPU NOT AVAILABLE')
    import sys; sys.exit(1)
"'''
out, err = run_full(c, gpu_test, timeout=60)
print(out)
if err:
    print(f"STDERR: {err}")
if 'ERROR: GPU NOT AVAILABLE' in out or 'CUDA available: False' in out:
    print("\n*** FATAL: PyTorch cannot see the GPU! ***")
    c.close()
    sys.exit(1)
print("\nGPU verified OK!")

# ============================================================
# PHASE 3: Start ComfyUI with correct GPU flags
# ============================================================
print("\n" + "="*60)
print("PHASE 3: Start ComfyUI with GPU")
print("="*60)

# The key insight: --novram was offloading EVERYTHING to CPU (0 MB on GPU)
# For this APU with shared memory, --lowvram is better:
# it keeps compute on GPU but swaps model layers in/out
# We also use --force-fp16 to reduce memory pressure
# --cpu-vae to avoid the known VAE decode hang on this GPU

COMFYUI_CMD = (
    f'{GPU_ENV} '
    'export OMP_NUM_THREADS=12; '
    'export MKL_NUM_THREADS=12; '
    'export OPENBLAS_NUM_THREADS=12; '
    'export MIOPEN_FIND_MODE=1; '  # Fast MIOpen kernel search
    'cd ~/ComfyUI && '
    'source ~/comfyui-env/bin/activate 2>/dev/null; . ~/comfyui-env/bin/activate 2>/dev/null; '
    'nohup python3 main.py '
    '--listen 0.0.0.0 --port 8188 '
    '--lowvram '
    '--force-fp16 '
    '--cpu-vae '
    '--disable-smart-memory '
    '> /tmp/comfyui.log 2>&1 &'
)

# Truncate old log first
run(c, 'truncate -s 0 /tmp/comfyui.log 2>/dev/null; touch /tmp/comfyui.log')
print("Starting ComfyUI with: --lowvram --force-fp16 --cpu-vae --disable-smart-memory")
print("(--lowvram keeps compute on GPU, swaps layers; --novram was wrong - it put everything on CPU)")
run(c, COMFYUI_CMD)
time.sleep(3)

# Verify it started
pid = run(c, 'pgrep -f "python.*main.py" 2>/dev/null')
if not pid:
    print("ERROR: ComfyUI failed to start!")
    log = run(c, 'cat /tmp/comfyui.log')
    print(f"Log:\n{log}")
    c.close()
    sys.exit(1)
print(f"ComfyUI started, PID: {pid}")

# Wait for server ready
print("Waiting for server ready...")
for i in range(60):
    try:
        resp = run(c, 'curl -s -o /dev/null -w "%{http_code}" http://127.0.0.1:8188/ 2>/dev/null', timeout=5)
        if resp == '200':
            print(f"Server ready after {i*3}s!")
            break
    except:
        pass
    # Also check for crash
    log_tail = run(c, 'tail -3 /tmp/comfyui.log 2>/dev/null')
    if 'Traceback' in log_tail or 'Error' in log_tail:
        print(f"Server log issue: {log_tail}")
    if i % 5 == 0 and i > 0:
        print(f"  [{i*3}s] Still waiting... log: {log_tail[-80:]}")
    time.sleep(3)
else:
    print("TIMEOUT waiting for ComfyUI!")
    log = run(c, 'tail -30 /tmp/comfyui.log')
    print(f"Log:\n{log}")
    c.close()
    sys.exit(1)

# Print startup log to confirm flags
log = run(c, 'head -20 /tmp/comfyui.log')
print(f"\nStartup log:\n{log}")

# ============================================================
# PHASE 4: Submit workflow via SFTP
# ============================================================
print("\n" + "="*60)
print("PHASE 4: Submit workflow")
print("="*60)

workflow = {
    "prompt": {
        "1": {
            "class_type": "UnetLoaderGGUF",
            "inputs": {
                "unet_name": "z_image_turbo-Q5_K_S.gguf"
            }
        },
        "2": {
            "class_type": "CLIPLoaderGGUF",
            "inputs": {
                "clip_name": "Qwen3-4B.i1-Q5_K_S.gguf",
                "type": "qwen_image"
            }
        },
        "3": {
            "class_type": "VAELoader",
            "inputs": {
                "vae_name": "ae.safetensors"
            }
        },
        "4": {
            "class_type": "CLIPTextEncode",
            "inputs": {
                "text": "A red fox in a snowy forest, photorealistic, highly detailed",
                "clip": ["2", 0]
            }
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {
                "width": 512,
                "height": 512,
                "batch_size": 1
            }
        },
        "6": {
            "class_type": "KSampler",
            "inputs": {
                "model": ["1", 0],
                "positive": ["4", 0],
                "negative": ["4", 0],
                "latent_image": ["5", 0],
                "seed": 42,
                "steps": 8,
                "cfg": 1.0,
                "sampler_name": "euler",
                "scheduler": "simple",
                "denoise": 1.0
            }
        },
        "7": {
            "class_type": "VAEDecode",
            "inputs": {
                "samples": ["6", 0],
                "vae": ["3", 0]
            }
        },
        "8": {
            "class_type": "SaveImage",
            "inputs": {
                "images": ["7", 0],
                "filename_prefix": "ZImageTurbo_GPU"
            }
        }
    }
}

# Write via SFTP
sftp = c.open_sftp()
wf_json = json.dumps(workflow)
with sftp.open('/tmp/wf.json', 'w') as f:
    f.write(wf_json)
sftp.close()
print("Workflow written to /tmp/wf.json via SFTP")

# Verify JSON
verify = run(c, 'python3 -c "import json; d=json.load(open(\'/tmp/wf.json\')); print(f\'Nodes: {list(d[chr(34)+chr(34) if False else \"prompt\"].keys())}\')"')
print(f"Verify: {verify}")

# Submit
resp = run(c, 'curl -s -X POST http://127.0.0.1:8188/prompt -H "Content-Type: application/json" -d @/tmp/wf.json 2>/dev/null')
print(f"Submit response: {resp}")

if 'error' in resp.lower() and 'prompt_id' not in resp.lower():
    print(f"\n*** SUBMISSION ERROR ***")
    # Check what went wrong
    log = run(c, 'tail -10 /tmp/comfyui.log')
    print(f"Log: {log}")
    c.close()
    sys.exit(1)

try:
    resp_data = json.loads(resp)
    prompt_id = resp_data.get('prompt_id', 'unknown')
    print(f"Prompt ID: {prompt_id}")
except:
    print("Could not parse response, continuing anyway...")

# ============================================================
# PHASE 5: Monitor generation with GPU tracking
# ============================================================
print("\n" + "="*60)
print("PHASE 5: Monitor generation (GPU must be active!)")
print("="*60)

start_time = time.time()
last_log_len = 0

for i in range(120):  # Up to 30 minutes
    elapsed = int(time.time() - start_time)
    
    # GPU metrics
    gpu_pct = run(c, 'cat /sys/class/drm/card0/device/gpu_busy_percent 2>/dev/null')
    gpu_temp = run(c, 'cat /sys/class/drm/card0/device/hwmon/hwmon*/temp1_input 2>/dev/null')
    temp_c = int(gpu_temp) // 1000 if gpu_temp.isdigit() else '?'
    
    # GPU power
    gpu_power = run(c, f'{GPU_ENV} rocm-smi -P 2>&1 | grep "Graphics Package" | grep -oP "[\\d.]+" | head -1')
    
    # Process info
    proc = run(c, 'ps -p $(pgrep -f "python.*main.py" | head -1) -o %cpu,%mem,rss --no-headers 2>/dev/null')
    
    # Log tail
    log = run(c, 'tail -5 /tmp/comfyui.log 2>/dev/null')
    last_line = log.split('\n')[-1] if log else ''
    
    # Output files
    files = run(c, 'ls ~/ComfyUI/output/*.png 2>/dev/null')
    
    # Queue
    queue = run(c, 'curl -s http://127.0.0.1:8188/queue 2>/dev/null')
    
    status = f"[{elapsed:>4}s] GPU:{gpu_pct:>3}% {temp_c}C {gpu_power}W | proc:{proc} | {last_line[-100:]}"
    print(status)
    
    # SUCCESS: Image generated!
    if files:
        print(f"\n{'='*60}")
        print(f"*** SUCCESS! IMAGE GENERATED! ***")
        print(f"Files: {files}")
        print(f"Total time: {elapsed}s")
        print(f"{'='*60}")
        
        # Print final log
        final_log = run(c, 'tail -20 /tmp/comfyui.log 2>/dev/null')
        print(f"\nFinal log:\n{final_log}")
        break
    
    # Check if queue is empty (job done or failed)
    try:
        qdata = json.loads(queue)
        running = len(qdata.get('queue_running', []))
        pending = len(qdata.get('queue_pending', []))
        if running == 0 and pending == 0 and elapsed > 30:
            print(f"\nQueue empty after {elapsed}s. Checking if image was saved...")
            time.sleep(2)
            files = run(c, 'ls ~/ComfyUI/output/*.png 2>/dev/null')
            if files:
                print(f"*** SUCCESS! {files}")
            else:
                print("No image. Checking log for errors:")
                err_log = run(c, 'tail -30 /tmp/comfyui.log 2>/dev/null')
                print(err_log)
            break
    except:
        pass
    
    # Check for process death
    alive = run(c, 'pgrep -f "python.*main.py" 2>/dev/null')
    if not alive:
        print("\n*** ComfyUI process died! ***")
        crash_log = run(c, 'tail -40 /tmp/comfyui.log 2>/dev/null')
        print(f"Crash log:\n{crash_log}")
        break
    
    time.sleep(15)

c.close()
print("\nDone.")