#!/usr/bin/env python3
"""
BC-250 v3 amdgpu kernel module build script.
Downloads kernel source, applies v3 patches, builds the module.
Runs entirely over SSH on the BC250.
"""
import paramiko
import time
import sys

KERNEL_VER = "6.19.6"
KERNEL_FULL = "6.19.6-2-cachyos"
BUILD_DIR = "/home/fabian/kernel-build"
SRC_DIR = f"{BUILD_DIR}/linux-{KERNEL_VER}"
AMDGPU_DIR = f"{SRC_DIR}/drivers/gpu/drm/amd/amdgpu"

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="", timeout=300):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    stdin, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode().strip()
    err = stderr.read().decode().strip()
    rc = stdout.channel.recv_exit_status()
    if out:
        # Truncate very long output
        lines = out.split('\n')
        if len(lines) > 50:
            print('\n'.join(lines[:20]))
            print(f"  ... ({len(lines)-40} lines omitted) ...")
            print('\n'.join(lines[-20:]))
        else:
            print(out)
    if err and rc != 0:
        print(f"STDERR: {err[:500]}")
    if rc != 0:
        print(f"EXIT CODE: {rc}")
    return out, rc

# ============================================================
# Step 1: Download kernel source
# ============================================================
out, rc = run(f"test -d {SRC_DIR} && echo EXISTS || echo MISSING")
if "EXISTS" in out:
    print(f"\nKernel source already exists at {SRC_DIR}")
else:
    run(f"mkdir -p {BUILD_DIR}", "Creating build directory")
    
    # Download kernel source tarball
    tarball_url = f"https://cdn.kernel.org/pub/linux/kernel/v6.x/linux-{KERNEL_VER}.tar.xz"
    run(f"cd {BUILD_DIR} && curl -LO {tarball_url}", 
        f"Downloading linux-{KERNEL_VER}.tar.xz from kernel.org", timeout=600)
    
    # Extract
    run(f"cd {BUILD_DIR} && tar xf linux-{KERNEL_VER}.tar.xz",
        "Extracting kernel source", timeout=300)
    
    # Clean up tarball
    run(f"rm {BUILD_DIR}/linux-{KERNEL_VER}.tar.xz")

# ============================================================
# Step 2: Prepare build environment
# ============================================================
run(f"""cd {SRC_DIR} && \\
  cp /usr/lib/modules/{KERNEL_FULL}/build/.config . && \\
  cp /usr/lib/modules/{KERNEL_FULL}/build/Module.symvers . && \\
  cp /usr/lib/modules/{KERNEL_FULL}/build/localversion.10-pkgrel . && \\
  cp /usr/lib/modules/{KERNEL_FULL}/build/localversion.20-pkgname . && \\
  echo 'Build files copied'""",
    "Copying kernel config, symvers, and localversion files")

# Prepare the kernel tree (generate required headers)
run(f"cd {SRC_DIR} && make LLVM=1 olddefconfig 2>&1 | tail -5",
    "Running olddefconfig", timeout=120)
run(f"cd {SRC_DIR} && make LLVM=1 modules_prepare 2>&1 | tail -10",
    "Preparing modules build", timeout=120)

# ============================================================
# Step 3: Verify the source files exist and find patch targets
# ============================================================
print("\n" + "="*60)
print("  Verifying source files")
print("="*60)

for f in ["gfx_v10_0.c", "gmc_v10_0.c", "amdgpu_gmc.c"]:
    out, rc = run(f"wc -l {AMDGPU_DIR}/{f}")
    print(f"  {f}: {out.split()[0]} lines")

# ============================================================
# Step 4: Apply v3 patches
# ============================================================

# --- Patch 1: gfx_v10_0.c - Disable GFXOFF for Cyan Skillfish ---
print("\n" + "="*60)
print("  Patch 1: gfx_v10_0.c — GFXOFF Disable")
print("="*60)

# Find the exact function and add our case
# Look for the switch statement in gfx_v10_0_check_gfxoff_flag
out, rc = run(f"grep -n 'case IP_VERSION(10, 1, 10)' {AMDGPU_DIR}/gfx_v10_0.c")
if rc != 0:
    print("ERROR: Could not find IP_VERSION(10,1,10) in gfx_v10_0.c!")
    sys.exit(1)

# Check if already patched
out, rc = run(f"grep -c 'BC-250' {AMDGPU_DIR}/gfx_v10_0.c")
if out.strip() != "0":
    print("Already patched, skipping.")
else:
    # The patch: add a case for IP_VERSION(10, 1, 3) after the existing default: break;
    # We need to find the closing "default:" case in gfx_v10_0_check_gfxoff_flag
    # and insert our case before it
    
    patch1_script = r"""
import re

filepath = '""" + AMDGPU_DIR + r"""/gfx_v10_0.c'
with open(filepath, 'r') as f:
    content = f.read()

# Find the pattern: after the IP_VERSION(10,1,10) case block, before 'default:'
# in gfx_v10_0_check_gfxoff_flag function
old_pattern = '\tdefault:\n\t\tbreak;\n\t}\n}'
# Find it specifically near gfx_v10_0_check_gfxoff_flag
func_start = content.find('static void gfx_v10_0_check_gfxoff_flag')
if func_start == -1:
    print("ERROR: function not found")
    exit(1)

# Find the 'default: break; } }' within this function
func_region = content[func_start:func_start+2000]
default_pos = func_region.find('\tdefault:\n\t\tbreak;\n\t}\n}')
if default_pos == -1:
    # Try different whitespace
    default_pos = func_region.find('default:\n')
    if default_pos == -1:
        print("ERROR: default case not found")
        print("Function region:\n" + func_region[:500])
        exit(1)
    # Get a bit more context
    context = func_region[default_pos-100:default_pos+100]
    print(f"Found default at offset {default_pos}, context:\n{context}")
    exit(1)

insert_pos = func_start + default_pos

new_case = '''\t/* ===== BC-250 v3 PATCH START ===== */
\tcase IP_VERSION(10, 1, 3):
\t\t/*
\t\t * BC-250 / Cyan Skillfish (gfx1013): GFXOFF causes the GPU to
\t\t * enter a power-saving state from which it cannot reliably wake.
\t\t * Unconditionally disable GFXOFF to prevent GPU hangs.
\t\t */
\t\tadev->pm.pp_feature &= ~PP_GFXOFF_MASK;
\t\tdev_info(adev->dev,
\t\t\t "BC-250: GFXOFF disabled to prevent GPU power-state hangs\\n");
\t\tbreak;
\t/* ===== BC-250 v3 PATCH END ===== */
'''

content = content[:insert_pos] + new_case + content[insert_pos:]

with open(filepath, 'w') as f:
    f.write(content)

print("Patch 1 applied successfully")
"""
    
    # Write patch script to remote
    with sftp.open('/tmp/patch1.py', 'w') as f:
        f.write(patch1_script)
    run("python3 /tmp/patch1.py", "Applying Patch 1")
    run(f"grep -A15 'gfx_v10_0_check_gfxoff_flag' {AMDGPU_DIR}/gfx_v10_0.c | head -30",
        "Verify Patch 1")

# --- Patch 2: gmc_v10_0.c - KIQ bypass + Dead-GPU detection ---
print("\n" + "="*60)
print("  Patch 2: gmc_v10_0.c — KIQ bypass + Dead-GPU")
print("="*60)

out, rc = run(f"grep -c 'BC-250' {AMDGPU_DIR}/gmc_v10_0.c")
if out.strip() != "0":
    print("Already patched, skipping.")
else:
    patch2_script = r"""
filepath = '""" + AMDGPU_DIR + r"""/gmc_v10_0.c'
with open(filepath, 'r') as f:
    content = f.read()

# ===== Patch 2a: KIQ bypass in gmc_v10_0_flush_gpu_tlb =====
# Find the line: if (adev->gfx.kiq[0].ring.sched.ready && !adev->enable_mes &&
# Insert KIQ bypass BEFORE it

kiq_check = 'if (adev->gfx.kiq[0].ring.sched.ready && !adev->enable_mes &&'
pos = content.find(kiq_check)
if pos == -1:
    # Try alternate: might use kiq_inst or different formatting
    kiq_check = 'if (adev->gfx.kiq['
    pos = content.find(kiq_check)
    if pos == -1:
        print("ERROR: Could not find KIQ check in gmc_v10_0_flush_gpu_tlb")
        exit(1)

# Make sure we're in gmc_v10_0_flush_gpu_tlb
func_start = content.rfind('gmc_v10_0_flush_gpu_tlb', 0, pos)
if func_start == -1:
    print("ERROR: Not in gmc_v10_0_flush_gpu_tlb?")
    exit(1)

# Find the start of the line (go back to newline)
line_start = content.rfind('\n', 0, pos) + 1
# Get indentation
indent = ''
for ch in content[line_start:pos]:
    if ch in ' \t':
        indent += ch
    else:
        break

kiq_bypass = indent + """/* ===== BC-250 v2 PATCH: KIQ bypass ===== */
""" + indent + """/* BC-250 / Cyan Skillfish (gfx1013): KIQ ring TLB flush hangs this GPU.
""" + indent + """ * Skip to direct MMIO register path. Widen to all gfx10.1.x for safety.
""" + indent + """ */
""" + indent + """{
""" + indent + """\tuint32_t gc_ver = amdgpu_ip_version(adev, GC_HWIP, 0);
""" + indent + """\tif ((gc_ver >= IP_VERSION(10, 1, 0)) && (gc_ver < IP_VERSION(10, 2, 0)))
""" + indent + """\t\tgoto use_mmio;
""" + indent + """}
""" + indent + """/* ===== BC-250 v2 PATCH END ===== */
"""

content = content[:line_start] + kiq_bypass + content[line_start:]

# ===== Patch 2b: Pre-spinlock health check before MMIO =====
# Find: "hub_ip = (vmhub == AMDGPU_GFXHUB(0))" in same function
# This is the start of the MMIO path ("use_mmio:" label or the code after it)

use_mmio_label = content.find('use_mmio:', content.find('gmc_v10_0_flush_gpu_tlb'))
if use_mmio_label == -1:
    # The label might not exist yet - we need to find where the MMIO path starts
    # In vanilla kernel it should be after "/* This path is needed before KIQ/MES/GFXOFF"
    mmio_comment = content.find('This path is needed before KIQ')
    if mmio_comment == -1:
        print("WARNING: Could not find MMIO path - label may already exist from bypass")

# Find "hub_ip = (vmhub == AMDGPU_GFXHUB(0))" after our inserted code
hub_ip_line = 'hub_ip = (vmhub == AMDGPU_GFXHUB(0))'
hub_pos = content.find(hub_ip_line, content.find('gmc_v10_0_flush_gpu_tlb'))
if hub_pos == -1:
    print("ERROR: Could not find hub_ip assignment")
    exit(1)

# If there's no use_mmio label, add one before the hub_ip line
if content.find('use_mmio:', content.find('gmc_v10_0_flush_gpu_tlb')) == -1:
    # Find the line start before hub_ip
    hub_line_start = content.rfind('\n', 0, hub_pos) + 1
    content = content[:hub_line_start] + "use_mmio:\n" + content[hub_line_start:]

# Re-find hub_ip position after possible label insertion
hub_pos = content.find(hub_ip_line, content.find('gmc_v10_0_flush_gpu_tlb'))
hub_line_end = content.find('\n', hub_pos)

# Insert health check AFTER hub_ip assignment line
health_check = """

\t/* ===== BC-250 v3 PATCH: Pre-spinlock health check ===== */
\t/*
\t * BC-250: GPU health check before entering the spinlock-protected
\t * MMIO section. On this SoC the internal PCIe fabric has NO completion
\t * timeout - readl() on an unresponsive GPU hangs CPU indefinitely.
\t */
\t{
\t\tuint32_t gc_ver = amdgpu_ip_version(adev, GC_HWIP, 0);

\t\tif ((gc_ver >= IP_VERSION(10, 1, 0)) &&
\t\t    (gc_ver < IP_VERSION(10, 2, 0))) {
\t\t\ttmp = RREG32_RLC_NO_KIQ(ack, hub_ip);
\t\t\tif (tmp == 0xFFFFFFFF) {
\t\t\t\tdev_err_ratelimited(adev->dev,
\t\t\t\t\t"BC-250: GPU unreachable (MMIO returned 0xFFFFFFFF), "
\t\t\t\t\t"skipping TLB flush vmid=%u hub=%u\\n",
\t\t\t\t\tvmid, vmhub);
\t\t\t\treturn;
\t\t\t}
\t\t}
\t}
\t/* ===== BC-250 v3 PATCH END ===== */
"""
content = content[:hub_line_end] + health_check + content[hub_line_end:]

# ===== Patch 2c: In-spinlock semaphore dead-GPU check =====
# Find the semaphore acquire loop: "if (tmp & 0x1)" inside the TLB flush function
# We need to add 0xFFFFFFFF check right after the RREG32 in the sem loop

# Find "a read return value of 1 means semaphore" comment (unique marker)
sem_comment = content.find('a read return value of 1 means semaphore', content.find('gmc_v10_0_flush_gpu_tlb'))
if sem_comment == -1:
    sem_comment = content.find('semaphore acq', content.find('gmc_v10_0_flush_gpu_tlb'))

if sem_comment != -1:
    # Find "if (tmp & 0x1)" after this comment
    tmp_check = content.find('if (tmp & 0x1)', sem_comment)
    if tmp_check != -1:
        # Insert dead-GPU check before "if (tmp & 0x1)"
        tmp_line_start = content.rfind('\n', 0, tmp_check) + 1
        sem_dead_check = """\t\t\t\t/* ===== BC-250 v3 PATCH: In-spinlock sem dead-GPU check ===== */
\t\t\t\tif (tmp == 0xFFFFFFFF) {
\t\t\t\t\tdev_err_ratelimited(adev->dev,
\t\t\t\t\t\t"BC-250: GPU died during sem acquire (0xFFFFFFFF)\\n");
\t\t\t\t\tspin_unlock(&adev->gmc.invalidate_lock);
\t\t\t\t\treturn;
\t\t\t\t}
\t\t\t\t/* ===== BC-250 v3 PATCH END ===== */
"""
        content = content[:tmp_line_start] + sem_dead_check + content[tmp_line_start:]
else:
    print("WARNING: Could not find semaphore comment - skipping sem dead-GPU check")

# ===== Patch 2d: ACK-wait loop dead-GPU check =====
# Find the ACK wait loop: "Wait for ACK with a delay" comment
ack_comment = content.find('Wait for ACK with a delay', content.find('gmc_v10_0_flush_gpu_tlb'))
if ack_comment != -1:
    # Find "tmp &= 1 << vmid;" after this - that's inside the ACK loop
    tmp_mask = content.find('tmp &= 1 << vmid', ack_comment)
    if tmp_mask != -1:
        tmp_mask_line_start = content.rfind('\n', 0, tmp_mask) + 1
        ack_dead_check = """\t\t\t\t/* ===== BC-250 v3 PATCH: ACK-wait dead-GPU check ===== */
\t\t\t\tif (tmp == 0xFFFFFFFF) {
\t\t\t\t\tdev_err_ratelimited(adev->dev,
\t\t\t\t\t\t"BC-250: GPU died during TLB flush ACK wait (0xFFFFFFFF)\\n");
\t\t\t\t\tif (use_semaphore)
\t\t\t\t\t\tWREG32_RLC_NO_KIQ(sem, 0, hub_ip);
\t\t\t\t\tspin_unlock(&adev->gmc.invalidate_lock);
\t\t\t\t\treturn;
\t\t\t\t}
\t\t\t\t/* ===== BC-250 v3 PATCH END ===== */
"""
        content = content[:tmp_mask_line_start] + ack_dead_check + content[tmp_mask_line_start:]
else:
    print("WARNING: Could not find ACK wait comment")

# ===== Patch 2e: gmc_v10_0_hw_init - PASID KIQ disable =====
# Find gmc_v10_0_hw_init function, specifically the line:
# adev->gmc.flush_pasid_uses_kiq = !amdgpu_emu_mode;
hw_init_func = content.find('static int gmc_v10_0_hw_init')
if hw_init_func == -1:
    print("ERROR: Could not find gmc_v10_0_hw_init")
    exit(1)

pasid_kiq = content.find('flush_pasid_uses_kiq', hw_init_func)
if pasid_kiq == -1:
    print("ERROR: Could not find flush_pasid_uses_kiq in hw_init")
    exit(1)

# Find the full line
pasid_line_start = content.rfind('\n', 0, pasid_kiq) + 1
pasid_line_end = content.find('\n', pasid_kiq)
old_line = content[pasid_line_start:pasid_line_end]

new_block = """\t/* ===== BC-250 v2 PATCH: Disable KIQ-based PASID flush ===== */
\t{
\t\tuint32_t gc_ver = amdgpu_ip_version(adev, GC_HWIP, 0);
\t\tif ((gc_ver >= IP_VERSION(10, 1, 0)) && (gc_ver < IP_VERSION(10, 2, 0)))
\t\t\tadev->gmc.flush_pasid_uses_kiq = false;
\t\telse
\t\t\tadev->gmc.flush_pasid_uses_kiq = !amdgpu_emu_mode;
\t}
\t/* ===== BC-250 v2 PATCH END ===== */"""

content = content[:pasid_line_start] + new_block + content[pasid_line_end:]

with open(filepath, 'w') as f:
    f.write(content)

print("Patch 2 applied successfully (5 sub-patches to gmc_v10_0.c)")
"""
    
    with sftp.open('/tmp/patch2.py', 'w') as f:
        f.write(patch2_script)
    run("python3 /tmp/patch2.py", "Applying Patch 2")
    run(f"grep -c 'BC-250' {AMDGPU_DIR}/gmc_v10_0.c", "Count BC-250 markers in gmc_v10_0.c")

# --- Patch 3: amdgpu_gmc.c - KIQ bypass + Dead-GPU detection ---
print("\n" + "="*60)
print("  Patch 3: amdgpu_gmc.c — KIQ bypass + Dead-GPU")
print("="*60)

out, rc = run(f"grep -c 'BC-250' {AMDGPU_DIR}/amdgpu_gmc.c")
if out.strip() != "0":
    print("Already patched, skipping.")
else:
    patch3_script = r"""
filepath = '""" + AMDGPU_DIR + r"""/amdgpu_gmc.c'
with open(filepath, 'r') as f:
    content = f.read()

# ===== Patch 3a: KIQ bypass in amdgpu_gmc_flush_gpu_tlb_pasid =====
# Find down_read_trylock check, then insert bypass after it

func_start = content.find('int amdgpu_gmc_flush_gpu_tlb_pasid')
if func_start == -1:
    print("ERROR: amdgpu_gmc_flush_gpu_tlb_pasid not found")
    exit(1)

# Find the trylock return 0 line
trylock = content.find('down_read_trylock', func_start)
if trylock == -1:
    print("ERROR: down_read_trylock not found")
    exit(1)

# Find the end of the if block (return 0;)
return_0 = content.find('return 0;', trylock)
end_of_block = content.find('\n', return_0) + 1

bypass_code = """
\t/* ===== BC-250 v2 PATCH: KIQ bypass for PASID flush ===== */
\t{
\t\tuint32_t gc_ver = amdgpu_ip_version(adev, GC_HWIP, 0);
\t\tpr_warn_once("amdgpu: flush_gpu_tlb_pasid called, GC_HWIP=0x%08x "
\t\t\t     "(10.1.3=0x%08x) kiq_flag=%d\\n",
\t\t\t     gc_ver, IP_VERSION(10, 1, 3),
\t\t\t     adev->gmc.flush_pasid_uses_kiq);

\t\tif ((gc_ver >= IP_VERSION(10, 1, 0)) && (gc_ver < IP_VERSION(10, 2, 0))) {
\t\t\tpr_warn_once("amdgpu: BC-250 KIQ bypass active "
\t\t\t\t     "(gc_ver=0x%08x)\\n", gc_ver);
\t\t\tadev->gmc.gmc_funcs->flush_gpu_tlb_pasid(adev, pasid,
\t\t\t\t\t\t\t\t  flush_type, all_hub,
\t\t\t\t\t\t\t\t  inst);
\t\t\tr = 0;
\t\t\tgoto error_unlock_reset;
\t\t}
\t}
\t/* ===== BC-250 v2 PATCH END ===== */
"""

content = content[:end_of_block] + bypass_code + content[end_of_block:]

# ===== Patch 3b: KIQ bypass + dead-GPU in amdgpu_gmc_fw_reg_write_reg_wait =====
func2_start = content.find('void amdgpu_gmc_fw_reg_write_reg_wait')
if func2_start == -1:
    print("ERROR: amdgpu_gmc_fw_reg_write_reg_wait not found")
    exit(1)

# Find the first KIQ-related code after the function signature
# Look for "ring->sched.ready" or "spin_lock"
# We need to insert BEFORE the KIQ ring submission code
# Find the first operational line after variable declarations

# Look for spin_lock_irqsave which starts the KIQ path
spinlock = content.find('spin_lock_irqsave', func2_start)
if spinlock == -1:
    # Try another marker
    spinlock = content.find('ring->sched.ready', func2_start)
    if spinlock == -1:
        print("ERROR: Could not find KIQ code in fw_reg_write_reg_wait")
        exit(1)

spinlock_line_start = content.rfind('\n', 0, spinlock) + 1

bypass2_code = """\t/* ===== BC-250 v2+v3 PATCH: KIQ bypass + dead-GPU detection ===== */
\t{
\t\tuint32_t gc_ver = amdgpu_ip_version(adev, GC_HWIP, 0);
\t\tif ((gc_ver >= IP_VERSION(10, 1, 0)) && (gc_ver < IP_VERSION(10, 2, 0))) {
\t\t\tuint32_t tmp;

\t\t\tpr_warn_once("amdgpu: BC-250 KIQ bypass active in "
\t\t\t\t     "fw_reg_write_reg_wait (gc=0x%08x)\\n", gc_ver);

\t\t\t/* v3: Health-check read before writing */
\t\t\ttmp = RREG32_NO_KIQ(reg1);
\t\t\tif (tmp == 0xFFFFFFFF) {
\t\t\t\tdev_err_ratelimited(adev->dev,
\t\t\t\t\t"BC-250: GPU unreachable in fw_reg_write_reg_wait "
\t\t\t\t\t"(reg1=0x%x returned 0xFFFFFFFF), skipping\\n", reg1);
\t\t\t\treturn;
\t\t\t}

\t\t\tWREG32_NO_KIQ(reg0, ref);
\t\t\tfor (cnt = 0; cnt < adev->usec_timeout; cnt++) {
\t\t\t\ttmp = RREG32_NO_KIQ(reg1);
\t\t\t\t/* v3: Dead-GPU detection in polling loop */
\t\t\t\tif (tmp == 0xFFFFFFFF) {
\t\t\t\t\tdev_err_ratelimited(adev->dev,
\t\t\t\t\t\t"BC-250: GPU died during reg_write_reg_wait "
\t\t\t\t\t\t"(0xFFFFFFFF at reg1=0x%x)\\n", reg1);
\t\t\t\t\treturn;
\t\t\t\t}
\t\t\t\tif ((tmp & mask) == (ref & mask))
\t\t\t\t\treturn;
\t\t\t\tudelay(1);
\t\t\t}
\t\t\tdev_warn(adev->dev, "BC-250: MMIO reg write/wait timeout "
\t\t\t\t "reg0=0x%x reg1=0x%x\\n", reg0, reg1);
\t\t\treturn;
\t\t}
\t}
\t/* ===== BC-250 v2+v3 PATCH END ===== */
"""

content = content[:spinlock_line_start] + bypass2_code + content[spinlock_line_start:]

with open(filepath, 'w') as f:
    f.write(content)

print("Patch 3 applied successfully (2 sub-patches to amdgpu_gmc.c)")
"""
    
    with sftp.open('/tmp/patch3.py', 'w') as f:
        f.write(patch3_script)
    run("python3 /tmp/patch3.py", "Applying Patch 3")
    run(f"grep -c 'BC-250' {AMDGPU_DIR}/amdgpu_gmc.c", "Count BC-250 markers in amdgpu_gmc.c")

# ============================================================
# Step 5: Verify all patches
# ============================================================
print("\n" + "="*60)
print("  Patch summary — BC-250 markers in all files")
print("="*60)
for f in ["gfx_v10_0.c", "gmc_v10_0.c", "amdgpu_gmc.c"]:
    run(f"grep -n 'BC-250' {AMDGPU_DIR}/{f}")

# ============================================================
# Step 6: Build the module
# ============================================================
print("\n" + "="*60)
print("  Building amdgpu module (LLVM=1)")
print("  This may take several minutes...")
print("="*60)

# Build using nohup + log file for long build
run(f"cd {SRC_DIR} && nohup make LLVM=1 -j$(nproc) M=drivers/gpu/drm/amd/amdgpu modules > /tmp/amdgpu-build.log 2>&1; echo BUILD_EXIT=$?",
    "Building amdgpu module", timeout=900)

# Check build result  
run("tail -30 /tmp/amdgpu-build.log", "Build log (last 30 lines)")

# Check if module was built
out, rc = run(f"ls -la {AMDGPU_DIR}/amdgpu.ko 2>&1")
if rc != 0:
    print("\nERROR: Module build failed!")
    run("grep -i error /tmp/amdgpu-build.log | head -20", "Build errors")
    sys.exit(1)

# ============================================================
# Step 7: Strip, compress, and install
# ============================================================
print("\n" + "="*60)
print("  Stripping and compressing module")
print("="*60)

run(f"ls -lh {AMDGPU_DIR}/amdgpu.ko", "Module size before strip")
run(f"strip --strip-debug {AMDGPU_DIR}/amdgpu.ko")
run(f"ls -lh {AMDGPU_DIR}/amdgpu.ko", "Module size after strip")
run(f"zstd -19 -f {AMDGPU_DIR}/amdgpu.ko", "Compressing with zstd-19", timeout=120)
run(f"ls -lh {AMDGPU_DIR}/amdgpu.ko.zst", "Compressed module size")

# Backup original module
MODULE_DIR = f"/usr/lib/modules/{KERNEL_FULL}/kernel/drivers/gpu/drm/amd/amdgpu"
run(f"sudo cp {MODULE_DIR}/amdgpu.ko.zst {MODULE_DIR}/amdgpu.ko.zst.original 2>/dev/null; echo done",
    "Backing up original module")

# Install new module
run(f"sudo cp {AMDGPU_DIR}/amdgpu.ko.zst {MODULE_DIR}/amdgpu.ko.zst",
    "Installing v3 patched module")

# Update module dependencies
run("sudo depmod -a", "Updating module dependencies")

# Verify installed module has BC-250 strings
run(f"zstd -d -c {MODULE_DIR}/amdgpu.ko.zst | strings | grep 'BC-250'",
    "Verifying BC-250 strings in installed module")

print("\n" + "="*60)
print("  BUILD AND INSTALL COMPLETE")
print("  Reboot required to load the v3 patched module.")
print("="*60)

# Clean up temp files
run("rm -f /tmp/patch1.py /tmp/patch2.py /tmp/patch3.py")

sftp.close()
ssh.close()
