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

Part of Layer 2 of the v3 three-layer protection. Modifies the GPU-generation-specific
TLB flush code to:
  a) Bypass KIQ ring for all gfx10.1.x → goto use_mmio
  b) Add pre-spinlock 0xFFFFFFFF health check
  c) Add dead-GPU detection in semaphore acquire loop
  d) Add dead-GPU detection in ACK-wait loop
  e) Disable KIQ-based PASID flush in hw_init

Usage:
    python3 patch2_gmc.py /path/to/drivers/gpu/drm/amd/amdgpu
"""
import sys
import os

if len(sys.argv) < 2:
    print("Usage: python3 patch2_gmc.py <amdgpu_source_dir>")
    print("  e.g.: python3 patch2_gmc.py ~/kernel-build/linux-6.19.6/drivers/gpu/drm/amd/amdgpu")
    sys.exit(1)

AMDGPU = sys.argv[1]
filepath = os.path.join(AMDGPU, "gmc_v10_0.c")

if not os.path.isfile(filepath):
    print(f"ERROR: File not found: {filepath}")
    sys.exit(1)

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:
    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:]

# ========================================
# Add '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:
    hub_ip_assign = content.find('hub_ip =', flush_func)
    if hub_ip_assign == -1:
        print("ERROR: hub_ip assignment not found")
        sys.exit(1)

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:
    content = content[:hub_line_start] + 'use_mmio:\n' + content[hub_line_start:]

# ========================================
# Patch 2b: Pre-spinlock health check after hub_ip assignment
# ========================================
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 multi-line ternary
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
# ========================================
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 detection)")