import paramiko

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')
sftp = ssh.open_sftp()

def run(cmd, desc=""):
    if desc:
        print(f"\n=== {desc} ===")
    stdin, stdout, stderr = ssh.exec_command(cmd)
    out = stdout.read().decode().strip()
    err = stderr.read().decode().strip()
    rc = stdout.channel.recv_exit_status()
    if out: print(out)
    if err: print(f"STDERR: {err}")
    return out, rc

# 1. Fix fish config - read current, append ROCm block, write back
print("=== Updating fish config ===")
current = sftp.open('/home/fabian/.config/fish/config.fish', 'r').read().decode()

rocm_block = """
# === ROCm / HIP Configuration for AMD BC-250 (v3) ===
# GPU target override (gfx1013 -> gfx1010 compatible)
set -gx HSA_OVERRIDE_GFX_VERSION 10.1.0

# Device selection
set -gx HIP_VISIBLE_DEVICES 0

# ROCm path
set -gx ROCM_PATH /opt/rocm
set -gx PATH /opt/rocm/bin $PATH

# CRITICAL: Disable SDMA engine (HW bugs on gfx1013)
set -gx HSA_ENABLE_SDMA 0

# Disable profiling tools (stability)
set -gx HSA_TOOLS_LIB ""
set -gx HSA_TOOLS_REPORT_LOAD_FAILURE 0

# NOTE: Do NOT set HIP_LAUNCH_BLOCKING=1 or GPU_MAX_HW_QUEUES=1
# These severely hurt performance and are unnecessary with v3 kernel patches.
"""

if "HSA_OVERRIDE_GFX_VERSION" not in current:
    with sftp.open('/home/fabian/.config/fish/config.fish', 'w') as f:
        f.write(current + rocm_block)
    print("ROCm env vars appended to fish config")
else:
    print("Already configured, skipping")

run("cat ~/.config/fish/config.fish", "Verify fish config")

# 2. Fix amdgpu.conf - write to /tmp first, then sudo mv
print("\n=== Creating /etc/modprobe.d/amdgpu.conf ===")
amdgpu_conf = """# AMD BC-250 (Cyan Skillfish / gfx1013) - ROCm Stability Parameters
# noretry=0              - Allow page fault retry (critical for shared memory / APU)
# gpu_recovery=1         - Enable GPU recovery on timeout
# sched_hw_submission=2  - Limit concurrent HW submissions (prevent queue overload)
# ppfeaturemask=0xfff73ef7 - Disable GFXOFF (bit 15), SCLK_DEEP_SLEEP (bit 3),
#                            and ULV (bit 8) to prevent unrecoverable power states.
options amdgpu noretry=0 gpu_recovery=1 sched_hw_submission=2 ppfeaturemask=0xfff73ef7
"""

with sftp.open('/tmp/amdgpu.conf', 'w') as f:
    f.write(amdgpu_conf)

run("sudo cp /tmp/amdgpu.conf /etc/modprobe.d/amdgpu.conf && rm /tmp/amdgpu.conf")
run("cat /etc/modprobe.d/amdgpu.conf", "Verify amdgpu.conf")

# 3. Verify Limine config (already done by previous script)
run("cat /etc/default/limine", "Verify Limine config")

print("\n=== All configuration complete ===")
sftp.close()
ssh.close()