"""Fix SDPA patch for BC-250 gfx1010: always use manual SDPA on CUDA"""
import sys

path = '/home/fabian/ComfyUI/bc250_softmax_patch.py'
with open(path, 'r') as f:
    content = f.read()

# Fix 1: patched_sdpa - ALWAYS use manual SDPA on CUDA (no threshold)
old_sdpa = '''def patched_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
    S = key.size(-2)
    if query.is_cuda and S > SAFE_SOFTMAX_THRESHOLD:
        if not getattr(patched_sdpa, '_logged', False):
            logger.warning(f"[BC-250] Manual SDPA: Q={list(query.shape)}, S={S}")
            patched_sdpa._logged = True
        return _safe_sdpa(query, key, value, attn_mask=attn_mask,
                         dropout_p=dropout_p, is_causal=is_causal, scale=scale)
    return _original_sdpa(query, key, value, attn_mask=attn_mask,
                         dropout_p=dropout_p, is_causal=is_causal, scale=scale)'''

new_sdpa = '''def patched_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
    # gfx1010: ALWAYS use manual SDPA on CUDA — built-in math backend kernel hangs
    if query.is_cuda:
        if not getattr(patched_sdpa, '_logged', False):
            S = key.size(-2)
            logger.warning(f"[BC-250] Manual SDPA (ALWAYS): Q={list(query.shape)}, S={S}")
            patched_sdpa._logged = True
        return _safe_sdpa(query, key, value, attn_mask=attn_mask,
                         dropout_p=dropout_p, is_causal=is_causal, scale=scale)
    return _original_sdpa(query, key, value, attn_mask=attn_mask,
                         dropout_p=dropout_p, is_causal=is_causal, scale=scale)'''

if old_sdpa not in content:
    print("ERROR: patched_sdpa not found!")
    sys.exit(1)
content = content.replace(old_sdpa, new_sdpa)
print("OK: patched_sdpa → always manual on CUDA")

# Fix 2: _safe_sdpa - add sync after GPU matmul+softmax operations
old_safe = '''def _safe_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
    L, S = query.size(-2), key.size(-2)
    if scale is None:
        scale = query.size(-1) ** -0.5
    attn_weight = torch.matmul(query, key.transpose(-2, -1)) * scale
    if is_causal:
        causal_mask = torch.triu(torch.ones(L, S, dtype=torch.bool, device=query.device), diagonal=1)
        attn_weight = attn_weight.masked_fill(causal_mask, float('-inf'))
    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_weight = attn_weight.masked_fill(~attn_mask, float('-inf'))
        else:
            attn_weight = attn_weight + attn_mask
    attn_weight = _safe_softmax_impl(attn_weight, dim=-1)
    if dropout_p > 0.0:
        attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p)
    return torch.matmul(attn_weight, value)'''

new_safe = '''def _safe_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
    L, S = query.size(-2), key.size(-2)
    if scale is None:
        scale = query.size(-1) ** -0.5
    # gfx1010: sync between GPU ops to prevent kernel queue buildup + hang
    attn_weight = torch.matmul(query, key.transpose(-2, -1)) * scale
    if query.is_cuda:
        torch.cuda.synchronize()
    if is_causal:
        causal_mask = torch.triu(torch.ones(L, S, dtype=torch.bool, device=query.device), diagonal=1)
        attn_weight = attn_weight.masked_fill(causal_mask, float('-inf'))
    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_weight = attn_weight.masked_fill(~attn_mask, float('-inf'))
        else:
            attn_weight = attn_weight + attn_mask
    attn_weight = _safe_softmax_impl(attn_weight, dim=-1)
    if query.is_cuda:
        torch.cuda.synchronize()
    if dropout_p > 0.0:
        attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p)
    output = torch.matmul(attn_weight, value)
    if query.is_cuda:
        torch.cuda.synchronize()
    return output'''

if old_safe not in content:
    print("ERROR: _safe_sdpa not found!")
    sys.exit(1)
content = content.replace(old_safe, new_safe)
print("OK: _safe_sdpa → sync after GPU ops")

# Fix 3: Also patch softmax to ALWAYS use manual on CUDA (same reason)
old_softmax = '''def patched_softmax(input, dim=None, _stacklevel=3, dtype=None):
    if dim is None:
        dim = -1
    if dtype is not None:
        input = input.to(dtype)
    if input.is_cuda and input.shape[dim] > SAFE_SOFTMAX_THRESHOLD:
        if not getattr(patched_softmax, '_logged', False):
            logger.warning(f"[BC-250] Manual F.softmax: shape={list(input.shape)}, dim={dim}")
            patched_softmax._logged = True
        return _safe_softmax_impl(input, dim)
    return _original_softmax(input, dim=dim)'''

new_softmax = '''def patched_softmax(input, dim=None, _stacklevel=3, dtype=None):
    if dim is None:
        dim = -1
    if dtype is not None:
        input = input.to(dtype)
    # gfx1010: ALWAYS use manual softmax on CUDA — native kernel unreliable
    if input.is_cuda:
        if not getattr(patched_softmax, '_logged', False):
            logger.warning(f"[BC-250] Manual F.softmax (ALWAYS): shape={list(input.shape)}, dim={dim}")
            patched_softmax._logged = True
        return _safe_softmax_impl(input, dim)
    return _original_softmax(input, dim=dim)'''

if old_softmax not in content:
    print("WARNING: patched_softmax not found (may already be fixed)")
else:
    content = content.replace(old_softmax, new_softmax)
    print("OK: patched_softmax → always manual on CUDA")

# Fix 4: Same for tensor.softmax
old_tsm = '''def patched_tensor_softmax(self, dim=-1, dtype=None):
    if dtype is not None:
        self = self.to(dtype)
    if self.is_cuda and self.shape[dim] > SAFE_SOFTMAX_THRESHOLD:
        if not getattr(patched_tensor_softmax, '_logged', False):
            logger.warning(f"[BC-250] Manual softmax: shape={list(self.shape)}, dim={dim}")
            patched_tensor_softmax._logged = True
        return _safe_softmax_impl(self, dim)
    return _original_tensor_softmax(self, dim=dim)'''

new_tsm = '''def patched_tensor_softmax(self, dim=-1, dtype=None):
    if dtype is not None:
        self = self.to(dtype)
    # gfx1010: ALWAYS use manual softmax on CUDA
    if self.is_cuda:
        if not getattr(patched_tensor_softmax, '_logged', False):
            logger.warning(f"[BC-250] Manual softmax (ALWAYS): shape={list(self.shape)}, dim={dim}")
            patched_tensor_softmax._logged = True
        return _safe_softmax_impl(self, dim)
    return _original_tensor_softmax(self, dim=dim)'''

if old_tsm not in content:
    print("WARNING: patched_tensor_softmax not found (may already be fixed)")
else:
    content = content.replace(old_tsm, new_tsm)
    print("OK: patched_tensor_softmax → always manual on CUDA")

# Fix 5: Update version string
content = content.replace('v17 ready', 'v19 ready — ALL CUDA ops manual (no native kernels)')
content = content.replace('Comprehensive Monkey-Patch v17', 'Comprehensive Monkey-Patch v19')

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

print("\nALL PATCHES APPLIED — v19")