#!/usr/bin/env python3
"""Phase 3: Build v3 amdgpu module from CachyOS kernel source."""
import paramiko
import sys
import os

KERNEL_FULL = "6.19.6-2-cachyos"
BUILD_DIR = "/home/fabian/kernel-build"
SRC_DIR = f"{BUILD_DIR}/cachyos-6.19.6-1"
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

# Step 1: Setup build environment
run(f"cp /usr/lib/modules/{KERNEL_FULL}/build/.config {SRC_DIR}/",
    "Copying CachyOS .config")
run(f"cp /usr/lib/modules/{KERNEL_FULL}/build/Module.symvers {SRC_DIR}/",
    "Copying Module.symvers")
run(f"cp /usr/lib/modules/{KERNEL_FULL}/build/localversion.10-pkgrel {SRC_DIR}/",
    "Copying localversion.10-pkgrel")
run(f"cp /usr/lib/modules/{KERNEL_FULL}/build/localversion.20-pkgname {SRC_DIR}/",
    "Copying localversion.20-pkgname")

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",
    "Running modules_prepare", timeout=300)

# Step 2: Apply patches (upload and run)
# Copy patch scripts from vanilla tree patching (they use the same logic)
desktop = r'C:\Users\fabia\Desktop'

# Upload patch scripts - update them for new source dir
for fname in ['patch1_gfxoff.py', 'patch2_gmc.py', 'patch3_amdgpu_gmc.py']:
    local = os.path.join(desktop, fname)
    with open(local, 'r') as f:
        content = f.read()
    # Replace old path with new CachyOS path
    content = content.replace(
        '/home/fabian/kernel-build/linux-6.19.6',
        '/home/fabian/kernel-build/cachyos-6.19.6-1'
    )
    with sftp.open(f'/tmp/{fname}', 'w') as f:
        f.write(content)
    print(f"Uploaded {fname} (updated path)")

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
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 refs")

# Step 3: Build the module
run(f"cd {SRC_DIR} && make LLVM=1 -j$(nproc) M=drivers/gpu/drm/amd/amdgpu modules 2>&1 | tail -40",
    "Building amdgpu module from CachyOS source", timeout=900)

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

# Step 4: Strip, compress, install
run(f"ls -lh {AMDGPU_DIR}/amdgpu.ko", "Before strip")
run(f"strip --strip-debug {AMDGPU_DIR}/amdgpu.ko", "Stripping")
run(f"ls -lh {AMDGPU_DIR}/amdgpu.ko", "After strip")
run(f"zstd -19 -f {AMDGPU_DIR}/amdgpu.ko", "Compressing zstd-19", timeout=120)
run(f"ls -lh {AMDGPU_DIR}/amdgpu.ko.zst", "Compressed size")

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

# Rebuild initramfs
run("sudo limine-update 2>&1 | tail -10", "Rebuilding initramfs", timeout=120)

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

# Cleanup
run("rm -f /tmp/patch1_gfxoff.py /tmp/patch2_gmc.py /tmp/patch3_amdgpu_gmc.py")
run(f"rm -f {BUILD_DIR}/cachyos-6.19.6-1.tar.gz")

print("\n" + "="*60)
print("  v3 MODULE BUILT FROM CachyOS SOURCE AND INSTALLED")
print("  Reboot required.")
print("="*60)

sftp.close()
ssh.close()