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, desc=""):
    if desc:
        print(f"\n=== {desc} ===")
    print(f"$ {cmd}")
    stdin, stdout, stderr = ssh.exec_command(cmd)
    out = stdout.read().decode()
    err = stderr.read().decode()
    rc = stdout.channel.recv_exit_status()
    if out.strip():
        print(out.strip())
    if err.strip():
        print(f"STDERR: {err.strip()}")
    if rc != 0:
        print(f"EXIT CODE: {rc}")
    return out.strip(), err.strip(), rc

# 1. Append ROCm env vars to fish config
fish_env_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.
"""

# Check if already configured
out, _, _ = run("cat ~/.config/fish/config.fish")
if "HSA_OVERRIDE_GFX_VERSION" in out:
    print("\nFish config already has ROCm vars, skipping.")
else:
    # Write the block to a temp file and append
    run(f"cat >> ~/.config/fish/config.fish << 'FISHEOF'{fish_env_block}FISHEOF",
        "Appending ROCm env vars to fish config")
    run("cat ~/.config/fish/config.fish", "Verify fish config")

# 2. Create /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
"""

out, _, _ = run("cat /etc/modprobe.d/amdgpu.conf 2>/dev/null || echo 'NOT_FOUND'")
if "NOT_FOUND" in out or "ppfeaturemask" not in out:
    run(f"sudo tee /etc/modprobe.d/amdgpu.conf << 'MODEOF'{amdgpu_conf}MODEOF",
        "Creating /etc/modprobe.d/amdgpu.conf")
    run("cat /etc/modprobe.d/amdgpu.conf", "Verify amdgpu.conf")
else:
    print("\namdgpu.conf already configured, skipping.")

# 3. Update Limine boot parameters
run("cat /etc/default/limine", "Current Limine config")

# Read current config
out, _, _ = run("cat /etc/default/limine")
if "amdgpu.gpu_recovery" in out:
    print("\nLimine already has amdgpu boot params, skipping.")
else:
    # We need to add amdgpu params to the KERNEL_CMDLINE
    # The current line likely looks like:
    # KERNEL_CMDLINE[default]="quiet nowatchdog splash rw rootflags=subvol=/@ root=UUID=..."
    # We need to add params before rootflags or at end of quoted string
    
    # Use sed to insert amdgpu params before rootflags
    sed_cmd = r"""sudo sed -i 's|rootflags=subvol=/@|amdgpu.gpu_recovery=1 amdgpu.noretry=0 amdgpu.dc=0 amdgpu.lockup_timeout=120000 amdgpu.ppfeaturemask=0xfff73ef7 rootflags=subvol=/@|' /etc/default/limine"""
    run(sed_cmd, "Adding amdgpu boot params to Limine config")
    run("cat /etc/default/limine", "Verify Limine config")
    
    # Apply limine update
    run("sudo limine-update", "Applying Limine update")

print("\n=== Configuration complete ===")
ssh.close()