#!/usr/bin/env python3
"""Phase 2: Upload patches, apply them, build and install the amdgpu module."""
import paramiko
import sys
import os

KERNEL_VER = "6.19.6"
KERNEL_FULL = "6.19.6-2-cachyos"
SRC_DIR = f"/home/fabian/kernel-build/linux-{KERNEL_VER}"
AMDGPU_DIR = f"{SRC_DIR}/drivers/gpu/drm/amd/amdgpu"
MODULE_DIR = f"/usr/lib/modules/{KERNEL_FULL}/kernel/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=600):
    if desc:
        print(f"\n=== {desc} ===")
    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:
        lines = out.split('\n')
        if len(lines) > 80:
            print('\n'.join(lines[:30]))
            print(f"  ... ({len(lines)-60} lines omitted) ...")
            print('\n'.join(lines[-30:]))
        else:
            print(out)
    if err and rc != 0:
        print(f"STDERR: {err[:1000]}")
    if rc != 0:
        print(f"EXIT CODE: {rc}")
    return out, rc

# Upload patch scripts
desktop = r'C:\Users\fabia\Desktop'
for fname in ['patch1_gfxoff.py', 'patch2_gmc.py', 'patch3_amdgpu_gmc.py']:
    local = os.path.join(desktop, fname)
    remote = f'/tmp/{fname}'
    sftp.put(local, remote)
    print(f"Uploaded {fname}")

# Apply patches
run("python3 /tmp/patch1_gfxoff.py", "Patch 1: gfx_v10_0.c - GFXOFF disable")
run("python3 /tmp/patch2_gmc.py", "Patch 2: gmc_v10_0.c - KIQ bypass + dead-GPU")
run("python3 /tmp/patch3_amdgpu_gmc.py", "Patch 3: amdgpu_gmc.c - KIQ bypass + dead-GPU")

# Verify patches
print("\n=== Patch verification: BC-250 markers ===")
for f in ["gfx_v10_0.c", "gmc_v10_0.c", "amdgpu_gmc.c"]:
    out, _ = run(f"grep -c 'BC-250' {AMDGPU_DIR}/{f}")
    print(f"  {f}: {out} BC-250 references")

# Build the module
run(f"cd {SRC_DIR} && make LLVM=1 -j$(nproc) M=drivers/gpu/drm/amd/amdgpu modules 2>&1",
    "Building amdgpu module (this takes a few minutes)", timeout=900)

# Check if build succeeded
out, rc = run(f"test -f {AMDGPU_DIR}/amdgpu.ko && echo SUCCESS || echo FAILED")
if "FAILED" in out:
    print("\nERROR: Module build failed!")
    run(f"cd {SRC_DIR} && make LLVM=1 -j1 M=drivers/gpu/drm/amd/amdgpu modules 2>&1 | grep -i error | head -20",
        "Build errors")
    sys.exit(1)

# Strip and compress
run(f"ls -lh {AMDGPU_DIR}/amdgpu.ko", "Module size before strip")
run(f"strip --strip-debug {AMDGPU_DIR}/amdgpu.ko", "Stripping debug info")
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 and install
run(f"sudo cp {MODULE_DIR}/amdgpu.ko.zst {MODULE_DIR}/amdgpu.ko.zst.original",
    "Backing up stock module")
run(f"sudo cp {AMDGPU_DIR}/amdgpu.ko.zst {MODULE_DIR}/amdgpu.ko.zst",
    "Installing v3 patched module")
run("sudo depmod -a", "Updating module dependencies")

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

# Clean up
run("rm -f /tmp/patch1_gfxoff.py /tmp/patch2_gmc.py /tmp/patch3_amdgpu_gmc.py")

print("\n" + "="*60)
print("  v3 PATCHED MODULE BUILT AND INSTALLED")
print("  A reboot is required to load the new module.")
print("="*60)

sftp.close()
ssh.close()