// ASTCMipTailBugRepro.swift
//
// Minimal repro for an Apple GPU driver bug:
// For block-compressed (ASTC) textures uploaded to PRIVATE storage via the blit encoder,
// the sampler/texel-fetch mis-addresses the tail of a mip level when ALL of the following hold:
//   1. that mip level's width in texels is an exact multiple of the 16KB page width
//      (32 blocks; =128 texels for ASTC 4x4, 192 for 6x6, 256 for 8x8),
//   2. the texture width is NOT divisible by 2^level (a partial block column exists upstream),
//   3. that mip level is taller than one page (>32 block rows, i.e. >128 texels for 4x4).
//
// The mis-addressed region (everything past the level's first 128 texel rows for 4x4) reads
// invalid memory, which ASTC-decodes to opaque magenta (1,0,1), or aliases valid data
// (wrong image). Uploading the same texture to SHARED storage via -[MTLTexture replaceRegion:...]
// produces a layout that is addressed correctly (no corruption).
//
// Reproduced on: Apple M5 (macOS), iPhone 13 (A15). 100% deterministic for a given size.
//
// Build & run (macOS, Apple Silicon):
//   swiftc -O -o ASTCMipTailBugRepro ASTCMipTailBugRepro.swift && ./ASTCMipTailBugRepro
//
// Expected output: "514x1024 blit-uploaded mip2: CORRUPTED" while the 512-wide control and
// the replaceRegion variant are intact.

import Foundation
import Metal

let device = MTLCreateSystemDefaultDevice()!
let queue = device.makeCommandQueue()!

// A legal "void-extent" ASTC block that decodes to one solid RGBA color (any block size).
func solidASTCBlock(_ r: UInt16, _ g: UInt16, _ b: UInt16) -> [UInt8] {
    var blk = [UInt8](repeating: 0, count: 16)
    blk[0] = 0xfc; blk[1] = 0xfd
    for i in 2...7 { blk[i] = 0xff }
    blk[8] = UInt8(r & 0xff); blk[9] = UInt8(r >> 8)
    blk[10] = UInt8(g & 0xff); blk[11] = UInt8(g >> 8)
    blk[12] = UInt8(b & 0xff); blk[13] = UInt8(b >> 8)
    blk[14] = 0xff; blk[15] = 0xff
    return blk
}
let RED = solidASTCBlock(0xffff, 0, 0), GREEN = solidASTCBlock(0, 0xffff, 0)

let shaderSrc = """
#include 
using namespace metal;
struct VOut { float4 pos [[position]]; float2 uv; };
vertex VOut vrx(uint vid [[vertex_id]]) {
    float2 c = float2(vid & 1, vid >> 1);
    VOut o; o.pos = float4(c * 2 - 1, 0, 1); o.uv = float2(c.x, 1 - c.y); return o;
}
// Sample ONE mip level explicitly. NOTE: the sampler must use a real mipFilter
// (linear), otherwise level() is silently ignored.
fragment float4 frg(VOut in [[stage_in]], texture2d t [[texture(0)]],
                    sampler s [[sampler(0)]], constant float& lod [[buffer(0)]]) {
    return t.sample(s, in.uv, level(lod));
}
"""
let lib = try! device.makeLibrary(source: shaderSrc, options: nil)
let pdesc = MTLRenderPipelineDescriptor()
pdesc.vertexFunction = lib.makeFunction(name: "vrx")
pdesc.fragmentFunction = lib.makeFunction(name: "frg")
pdesc.colorAttachments[0].pixelFormat = .bgra8Unorm
let pso = try! device.makeRenderPipelineState(descriptor: pdesc)
let sdesc = MTLSamplerDescriptor()
sdesc.minFilter = .linear; sdesc.magFilter = .linear; sdesc.mipFilter = .linear
let sampler = device.makeSamplerState(descriptor: sdesc)!

// Build a WxH ASTC_4x4 texture (full mip chain), every level = top half red / bottom half green.
func makeTexture(_ W: Int, _ H: Int, useBlit: Bool) -> MTLTexture {
    var mips: [(Int, Int, Int, [UInt8])] = []
    var w = W, h = H
    let levels = Int(floor(log2(Double(max(W, H))))) + 1
    for _ in 0.. blit encoder -> private texture.
        var all = [UInt8](); for m in mips { all.append(contentsOf: m.3) }
        let buf = device.makeBuffer(bytes: all, length: all.count)!
        let cmd = queue.makeCommandBuffer()!; let blit = cmd.makeBlitCommandEncoder()!
        var off = 0
        for (i, m) in mips.enumerated() {
            blit.copy(from: buf, sourceOffset: off,
                      sourceBytesPerRow: m.2, sourceBytesPerImage: m.3.count,
                      sourceSize: MTLSize(width: m.0, height: m.1, depth: 1),
                      to: tex, destinationSlice: 0, destinationLevel: i, destinationOrigin: MTLOrigin())
            off += m.3.count
        }
        blit.endEncoding(); cmd.commit(); cmd.waitUntilCompleted()
    } else {
        // replaceRegion path (shared storage).
        for (i, m) in mips.enumerated() {
            m.3.withUnsafeBytes { ptr in
                tex.replace(region: MTLRegion(origin: .init(), size: MTLSize(width: m.0, height: m.1, depth: 1)),
                            mipmapLevel: i, withBytes: ptr.baseAddress!, bytesPerRow: m.2)
            }
        }
    }
    return tex
}

// Render one mip level 1:1 and check whether the bottom half still reads green.
func bottomHalfIntact(_ tex: MTLTexture, _ W: Int, _ H: Int, level: Int) -> Bool {
    let mw = max(1, W >> level), mh = max(1, H >> level)
    let rtd = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .bgra8Unorm, width: mw, height: mh, mipmapped: false)
    rtd.storageMode = .shared; rtd.usage = [.renderTarget]
    let rt = device.makeTexture(descriptor: rtd)!
    let rpd = MTLRenderPassDescriptor()
    rpd.colorAttachments[0].texture = rt
    rpd.colorAttachments[0].loadAction = .clear; rpd.colorAttachments[0].storeAction = .store
    let cmd = queue.makeCommandBuffer()!; let enc = cmd.makeRenderCommandEncoder(descriptor: rpd)!
    enc.setRenderPipelineState(pso)
    enc.setFragmentTexture(tex, index: 0)
    enc.setFragmentSamplerState(sampler, index: 0)
    var lod = Float(level)
    enc.setFragmentBytes(&lod, length: MemoryLayout.size, index: 0)
    enc.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
    enc.endEncoding(); cmd.commit(); cmd.waitUntilCompleted()
    var pix = [UInt8](repeating: 0, count: mw * mh * 4)
    rt.getBytes(&pix, bytesPerRow: mw * 4,
                from: MTLRegion(origin: .init(), size: MTLSize(width: mw, height: mh, depth: 1)), mipmapLevel: 0)
    for y in (mh / 2 + 2).. 200 && pix[i + 2] < 80) { return false }   // not green -> corrupted
    } }
    return true
}

func report(_ W: Int, _ H: Int, _ level: Int, useBlit: Bool, _ label: String) {
    let tex = makeTexture(W, H, useBlit: useBlit)
    let ok = bottomHalfIntact(tex, W, H, level: level)
    print("\(label): \(ok ? "INTACT" : "CORRUPTED (bottom of mip\(level) mis-addressed)")")
}

print("Device: \(device.name)")
report(514, 1024, 2, useBlit: true,  "514x1024  blit(private)  mip2")   // expected: CORRUPTED
report(512, 1024, 2, useBlit: true,  "512x1024  blit(private)  mip2")   // expected: INTACT (control)
report(514, 1024, 2, useBlit: false, "514x1024  replaceRegion  mip2")   // expected: INTACT (workaround)