#!/usr/bin/env python3
"""Step 2: Create venv and clone PyTorch on BC-250."""
import paramiko
import sys

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=600, 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():
        # Print last portion for long outputs
        lines = out.strip().split('\n')
        if len(lines) > 50:
            print(f"  ... ({len(lines)} lines, showing last 50)")
            print('\n'.join(lines[-50:]))
        else:
            print(out.strip())
    if err.strip():
        lines = [l for l in err.strip().split('\n') if 'warning:' not in l.lower()]
        if lines:
            if len(lines) > 30:
                print(f"STDERR ({len(lines)} lines, last 30):")
                print('\n'.join(lines[-30:]))
            else:
                print(f"STDERR: {chr(10).join(lines)}")
    print(f"  Exit code: {rc}")
    return rc, out, err

# Create the venv
run("python3 -m venv ~/comfyui-env --system-site-packages",
    desc="Create venv with system site-packages (for numpy, etc.)")

# Activate and install basic build deps in venv
run("bash -c 'source ~/comfyui-env/bin/activate && pip install --upgrade pip setuptools wheel'",
    desc="Upgrade pip in venv")

run("bash -c 'source ~/comfyui-env/bin/activate && pip install cmake ninja pyyaml typing-extensions cffi future six requests dataclasses filelock sympy networkx jinja2 numpy'",
    desc="Install PyTorch build deps in venv", timeout=120)

# Check if PyTorch source already exists
rc, out, _ = run("test -d ~/pytorch && echo EXISTS || echo MISSING",
    desc="Check for existing PyTorch source")

if "EXISTS" in out:
    print("\n  PyTorch source directory exists. Checking if it's a valid repo...")
    run("cd ~/pytorch && git log --oneline -1 2>&1", desc="Check PyTorch repo")
else:
    # Clone PyTorch — this is the big download
    print("\n  Cloning PyTorch (this will take a while)...")
    run("git clone --depth 1 --recursive --shallow-submodules https://github.com/pytorch/pytorch.git ~/pytorch 2>&1 | tail -20",
        desc="Clone PyTorch (shallow, with submodules)",
        timeout=1200)  # 20 minutes timeout

# Verify clone
run("ls -la ~/pytorch/setup.py ~/pytorch/torch/ 2>&1 | head -5",
    desc="Verify PyTorch source")

run("cd ~/pytorch && git log --oneline -1",
    desc="PyTorch version")

run("du -sh ~/pytorch",
    desc="PyTorch source size")

ssh.close()
print("\n\nDone — venv created and PyTorch cloned.")