#!/usr/bin/env python3
"""Patch 2: gmc_v10_0.c - KIQ bypass + Dead-GPU detection (5 sub-patches)"""
import sys

AMDGPU = "/home/fabian/kernel-build/linux-6.19.6/drivers/gpu/drm/amd/amdgpu"
filepath = f"{AMDGPU}/gmc_v10_0.c"

with open(filepath, 'r') as f:
    content = f.read()

if 'BC-250' in content:
    print("Already patched, skipping.")
    sys.exit(0)

# ========================================
# Patch 2a: KIQ bypass in gmc_v10_0_flush_gpu_tlb
# Insert BEFORE the KIQ check: if (adev->gfx.kiq[0].ring.sched.ready
# ========================================
flush_func = content.find('gmc_v10_0_flush_gpu_tlb(struct amdgpu_device')
if flush_func == -1:
    print("ERROR: gmc_v10_0_flush_gpu_tlb not found")
    sys.exit(1)

# Find the KIQ readiness check
kiq_check = content.find('adev->gfx.kiq[0].ring.sched.ready', flush_func)
if kiq_check == -1:
    # Try alternate patterns
    kiq_check = content.find('kiq[0].ring.sched.ready', flush_func)
    if kiq_check == -1:
        print("ERROR: KIQ readiness check not found in flush_gpu_tlb")
        sys.exit(1)

# Go back to the 'if' statement start
if_start = content.rfind('if (', flush_func, kiq_check)
if if_start == -1:
    if_start = kiq_check

# Find start of line
line_start = content.rfind('\n', 0, if_start) + 1

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

content = content[:line_start] + bypass + content[line_start:]

# ========================================
# Now add the 'use_mmio:' label before the MMIO path
# Find "hub_ip = (vmhub ==" which starts the MMIO code path
# ========================================
hub_ip_assign = content.find('hub_ip = (vmhub ==', flush_func)
if hub_ip_assign == -1:
    # In newer kernels it might be different
    hub_ip_assign = content.find('hub_ip =', flush_func)
    if hub_ip_assign == -1:
        print("ERROR: hub_ip assignment not found")
        sys.exit(1)

# Check if there's already a use_mmio label
hub_line_start = content.rfind('\n', 0, hub_ip_assign) + 1
preceding_text = content[hub_line_start - 50:hub_line_start].strip()
if 'use_mmio' not in preceding_text:
    # Find the comment before the MMIO path to place label after it
    mmio_comment = content.rfind('/*', flush_func, hub_ip_assign)
    comment_block_start = content.rfind('\n', 0, mmio_comment) + 1 if mmio_comment > flush_func else hub_line_start
    
    # Insert use_mmio label
    content = content[:hub_line_start] + 'use_mmio:\n' + content[hub_line_start:]

# ========================================
# Patch 2b: Pre-spinlock health check after hub_ip assignment  
# ========================================
# Re-find hub_ip
hub_ip_assign = content.find('hub_ip = (vmhub ==', flush_func)
if hub_ip_assign == -1:
    hub_ip_assign = content.find('hub_ip =', flush_func)
hub_line_end = content.find('\n', hub_ip_assign)

# Check for second line (MMHUB case) - the ternary might be split across 2 lines
next_line = content[hub_line_end+1:hub_line_end+100]
if next_line.strip().startswith(':') or next_line.strip().startswith('?'):
    hub_line_end = content.find('\n', hub_line_end + 1)

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

# ========================================
# Patch 2c: In-spinlock semaphore dead-GPU check
# Find: "if (tmp & 0x1)" inside the sem acquire loop
# ========================================
sem_marker = content.find('semaphore acq', flush_func)
if sem_marker == -1:
    sem_marker = content.find('a read return value of 1 means semaphore', flush_func)

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

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

# ========================================
# Patch 2e: gmc_v10_0_hw_init - PASID KIQ disable
# Replace: adev->gmc.flush_pasid_uses_kiq = !amdgpu_emu_mode;
# ========================================
hw_init = content.find('static int gmc_v10_0_hw_init')
if hw_init == -1:
    print("ERROR: gmc_v10_0_hw_init not found")
    sys.exit(1)

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

# Get the full line
pasid_line_start = content.rfind('\n', 0, pasid_kiq) + 1
pasid_line_end = content.find('\n', pasid_kiq)

new_block = (
    '\t/* ===== BC-250 v2 PATCH: Disable KIQ-based PASID flush ===== */\n'
    '\t{\n'
    '\t\tuint32_t gc_ver = amdgpu_ip_version(adev, GC_HWIP, 0);\n'
    '\t\tif ((gc_ver >= IP_VERSION(10, 1, 0)) && (gc_ver < IP_VERSION(10, 2, 0)))\n'
    '\t\t\tadev->gmc.flush_pasid_uses_kiq = false;\n'
    '\t\telse\n'
    '\t\t\tadev->gmc.flush_pasid_uses_kiq = !amdgpu_emu_mode;\n'
    '\t}\n'
    '\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: gmc_v10_0.c - 5 sub-patches (KIQ bypass + dead-GPU)")