content: {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  module Main where
  
  -- \|
  -- Example     : GPUCopyAndReadback
  -- Description : Tests GPU texture/buffer copy, blit, upload, download, and data integrity.
  -- Copyright   : (c) Kyle Lukaszek, 2025
  -- License     : BSD3
  --
  -- Based on SDL_gpu_examples/CopyAndReadback.c example.
  -- Demonstrates:
  -- - Uploading image and buffer data to the GPU.
  -- - GPU-to-GPU copies: `SDL_CopyGPUTextureToTexture`, `SDL_CopyGPUBufferToBuffer`.
  -- - GPU blit: `SDL_BlitGPUTexture` (including scaling).
  -- - GPU-to-CPU downloads: `SDL_DownloadFromGPUTexture`, `SDL_DownloadFromGPUBuffer`.
  -- - Synchronization with `SDL_GPUFence` for readback.
  -- - Verification of downloaded data against original data.
  -- - Displays original, copied, and blitted textures.
  -- \|
  
  import Control.Exception (bracket, bracketOnError, finally)
  import Control.Monad (unless, void, when, (>=>))
  import Data.Bits ((.|.))
  import qualified Data.ByteString as BS
  import qualified Data.ByteString.Unsafe as BSU
  import Data.IORef
  import Data.Maybe (isJust)
  import Data.Word (Word32)
  import Foreign.Marshal.Array (peekArray, pokeArray)
  import Foreign.Marshal.Utils (copyBytes)
  import Foreign.Ptr (castPtr, plusPtr)
  import Foreign.Storable (peek)
  import GPUCommon
  import SDL
  import System.Exit (exitFailure, exitSuccess)
  import System.FilePath ((</>))
  
  -- AppResources
  data AppResources = AppResources
    { resOriginalTexture :: SDLGPUTexture,
      resTextureCopy :: SDLGPUTexture,
      resTextureSmall :: SDLGPUTexture,
      resOriginalBuffer :: SDLGPUBuffer,
      resBufferCopy :: SDLGPUBuffer,
      -- Store original image and buffer data for comparison
      originalImageDataBytes :: BS.ByteString,
      originalBufferData :: [Word32],
      textureWidth :: Int,
      textureHeight :: Int
    }
  
  -- Test buffer data
  bufferData :: [Word32]
  bufferData = [2, 4, 8, 16, 32, 64, 128] -- 7 elements = 28 bytes
  
  -- main
  main :: IO ()
  main = do
    sdlLog $ "Compiled SDL Version: " ++ show sdlVersion
    linkedVersion <- sdlGetVersion
    sdlLog $ "Linked SDL Version: " ++ show linkedVersion
    maybeResult <- withContext "SDL3 Haskell GPU Copy & Readback" [] runAppGPU
    case maybeResult of
      Nothing -> sdlLog "Application initialization failed." >> exitFailure
      Just _ -> sdlLog "Application finished successfully." >> exitSuccess
  
  -- runAppGPU
  runAppGPU :: Context -> IO ()
  runAppGPU context = do
    sdlLog "Base context initialized."
    bracket
      (createAndTestResources context)
      (releaseResources context)
      $ \case
        Nothing -> sdlLog "Failed to create or test resources. Exiting."
        Just resources -> do
          sdlLog "Resources created, tested, and operations successful. Entering display loop."
          eventLoopGPU context resources
  
  -- createAndTestResources: Combines Init() and the initial test logic from C
  createAndTestResources :: Context -> IO (Maybe AppResources)
  createAndTestResources Context {..} = do
    sdlLog "--- Beginning Resource Creation and Initial Tests ---"
  
    -- Load the source image (bracket ensures surface is freed if loadImage fails or subsequent ops fail)
    bracketOnError
      (loadImage ("Content" </> "Images" </> "ravioli.bmp"))
      (\case Just surf -> sdlDestroySurface surf; _ -> pure ())
      $ \case
        Nothing -> sdlLog "!!! Failed to load source image." >> return Nothing
        Just surfacePtr ->
          -- This inner bracket is to ensure the loaded surfacePtr is freed
          -- after we're done with its pixel data, even if later steps fail.
          bracketOnError (pure surfacePtr) (\_ -> sdlLog "Releasing loaded surface from inner bracket" >> sdlDestroySurface surfacePtr) $ \surf -> do
            surfaceData <- peek surf
            let imgW = surfaceW surfaceData
            let imgH = surfaceH surfaceData
            let imgPitch = surfacePitch surfaceData
            let imgPixelsPtr = surfacePixels surfaceData
            let imgDataSize = imgH * imgPitch
            originalBS <- BSU.unsafePackCStringLen (castPtr imgPixelsPtr, imgDataSize)
  
            -- Create Textures & Buffers
            let commonTexCI = (defaultTextureCreateInfo imgW imgH) {texInfoUsage = SDL_GPU_TEXTUREUSAGE_SAMPLER}
            maybeOrigTex <- sdlCreateGPUTexture contextDevice commonTexCI
            maybeTexCopy <- sdlCreateGPUTexture contextDevice commonTexCI
            maybeTexSmall <-
              sdlCreateGPUTexture
                contextDevice
                ( (defaultTextureCreateInfo (imgW `div` 2) (imgH `div` 2))
                    { texInfoUsage = SDL_GPU_TEXTUREUSAGE_SAMPLER .|. SDL_GPU_TEXTUREUSAGE_COLOR_TARGET
                    }
                )
  
            (_, _, bufferDataSizeW32) <- calculateBufferDataSize bufferData "OriginalBufferData"
            let bufferUsage = SDL_GPU_BUFFERUSAGE_GRAPHICS_STORAGE_READ
            maybeOrigBuf <- createGPUBuffer contextDevice bufferUsage bufferDataSizeW32 "OriginalBuffer"
            maybeBufCopy <- createGPUBuffer contextDevice bufferUsage bufferDataSizeW32 "BufferCopy"
  
            let totalUploadSize = fromIntegral imgDataSize + bufferDataSizeW32
            let totalDownloadSize = totalUploadSize
            maybeUploadTransfer <- createTransferBuffer contextDevice totalUploadSize SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD "UploadCombined"
            maybeDownloadTransfer <- createTransferBuffer contextDevice totalDownloadSize SDL_GPU_TRANSFERBUFFERUSAGE_DOWNLOAD "DownloadCombined"
  
            -- Use a nested structure to handle Maybe resource creation and operations
            -- This ensures cleanup happens correctly for parts that were successfully created.
            case (maybeOrigTex, maybeTexCopy, maybeTexSmall, maybeOrigBuf, maybeBufCopy, maybeUploadTransfer, maybeDownloadTransfer) of
              (Just oTex, Just tCopy, Just tSmall, Just oBuf, Just bCopy, Just upTrans, Just downTrans) -> do
                mapAndPokeOk <- bracket
                  (sdlMapGPUTransferBuffer contextDevice upTrans False)
                  (\mptr -> when (isJust mptr) $ sdlUnmapGPUTransferBuffer contextDevice upTrans)
                  $ \case
                    Nothing -> sdlLog "Failed to map upload transfer buffer" >> return False
                    Just mappedUpPtr -> do
                      copyBytes mappedUpPtr imgPixelsPtr imgDataSize
                      pokeArray (castPtr mappedUpPtr `plusPtr` imgDataSize) bufferData
                      return True
  
                if not mapAndPokeOk
                  then do
                    sdlLog "Map and poke to upload buffer failed."
                    -- Release resources created so far in this 'Just' block before returning Nothing
                    sdlReleaseGPUTexture contextDevice oTex
                    sdlReleaseGPUTexture contextDevice tCopy
                    sdlReleaseGPUTexture contextDevice tSmall
                    sdlReleaseGPUBuffer contextDevice oBuf
                    sdlReleaseGPUBuffer contextDevice bCopy
                    sdlReleaseGPUTransferBuffer contextDevice upTrans
                    sdlReleaseGPUTransferBuffer contextDevice downTrans
                    return Nothing
                  else do
                    mCmdBuf <- sdlAcquireGPUCommandBuffer contextDevice
                    case mCmdBuf of
                      Nothing -> sdlLog "Failed to acquire command buffer" >> return Nothing -- Further cleanup needed if this was the only failure
                      Just cmdBuf -> do
                        -- Perform GPU operations, assuming success until a step fails
                        -- First Copy Pass
                        validAfterCopyPass1 <- do
                          -- Introduce a sub-do block to capture the Bool
                          mCopyPass1 <- sdlBeginGPUCopyPass cmdBuf
                          case mCopyPass1 of
                            Nothing -> sdlLog "Failed to begin first copy pass" >> return False
                            Just cp1 -> do
                              let upTexSrcInfo = SDLGPUTextureTransferInfo upTrans 0 (fromIntegral imgW) (fromIntegral imgH)
                              let upTexDstRegion = defaultTextureRegion oTex imgW imgH
                              sdlUploadToGPUTexture cp1 upTexSrcInfo upTexDstRegion False
                              let locO = SDLGPUTextureLocation oTex 0 0 0 0 0
                              let locC = SDLGPUTextureLocation tCopy 0 0 0 0 0
                              sdlCopyGPUTextureToTexture cp1 locO locC (fromIntegral imgW) (fromIntegral imgH) 1 False
                              let upBufSrcLoc = SDLGPUTransferBufferLocation upTrans (fromIntegral imgDataSize)
                              let upBufDstReg = SDLGPUBufferRegion oBuf 0 bufferDataSizeW32
                              sdlUploadToGPUBuffer cp1 upBufSrcLoc upBufDstReg False
                              let locOBuf = SDLGPUBufferLocation oBuf 0
                              let locCBuf = SDLGPUBufferLocation bCopy 0
                              sdlCopyGPUBufferToBuffer cp1 locOBuf locCBuf bufferDataSizeW32 False
                              sdlEndGPUCopyPass cp1
                              return True -- Indicate success of this block
                        if not validAfterCopyPass1
                          then
                            return Nothing -- Propagate failure for the whole createAndTestResources
                          else do
                            -- Blit Operation (Only if copyPass1 was successful)
                            let blitSrcReg = defaultBlitRegion oTex imgW imgH
                            let blitDstReg = defaultBlitRegion tSmall (imgW `div` 2) (imgH `div` 2)
                            let blitInfo = SDLGPUBlitInfo blitSrcReg blitDstReg SDL_GPU_LOADOP_DONT_CARE (SDLFColor 0 0 0 0) SDL_FLIP_NONE SDL_GPU_FILTER_LINEAR False
                            sdlBlitGPUTexture cmdBuf blitInfo
  
                            -- Second Copy Pass (Downloads)
                            mCopyPass2 <- sdlBeginGPUCopyPass cmdBuf
                            case mCopyPass2 of
                              Nothing -> sdlLog "Failed to begin second copy pass" >> return Nothing
                              Just cp2 -> do
                                let downTexSrcReg = defaultTextureRegion tCopy imgW imgH
                                let downTexDstInfo = SDLGPUTextureTransferInfo downTrans 0 (fromIntegral imgW) (fromIntegral imgH)
                                sdlDownloadFromGPUTexture cp2 downTexSrcReg downTexDstInfo
                                let downBufSrcReg = SDLGPUBufferRegion bCopy 0 bufferDataSizeW32
                                let downBufDstLoc = SDLGPUTransferBufferLocation downTrans (fromIntegral imgDataSize)
                                sdlDownloadFromGPUBuffer cp2 downBufSrcReg downBufDstLoc
                                sdlEndGPUCopyPass cp2
  
                                mFence <- sdlSubmitGPUCommandBufferAndAcquireFence cmdBuf
                                case mFence of
                                  Nothing -> sdlLog "Failed to acquire fence" >> return Nothing
                                  Just fence -> do
                                    waitOk <- sdlWaitForGPUFences contextDevice True [fence]
                                    sdlReleaseGPUFence contextDevice fence
  
                                    if waitOk
                                      then bracket
                                        (sdlMapGPUTransferBuffer contextDevice downTrans False)
                                        (\mptr -> when (isJust mptr) $ sdlUnmapGPUTransferBuffer contextDevice downTrans)
                                        $ \case
                                          Nothing -> sdlLog "Failed to map download buffer" >> return Nothing
                                          Just mappedDownPtr -> do
                                            downloadedTexBS <- BSU.unsafePackCStringLen (castPtr mappedDownPtr, imgDataSize)
                                            let texMatch = downloadedTexBS == originalBS
                                            if texMatch
                                              then sdlLog "SUCCESS! Texture data matches."
                                              else sdlLog "FAILURE! Texture data MISMATCH."
  
                                            let downloadedBufPtr = castPtr mappedDownPtr `plusPtr` imgDataSize
                                            downloadedBufData <- peekArray (length bufferData) downloadedBufPtr :: IO [Word32]
                                            let bufMatch = downloadedBufData == bufferData
                                            if bufMatch
                                              then sdlLog "SUCCESS! Buffer data matches."
                                              else sdlLog $ "FAILURE! Buffer data MISMATCH. Expected: " ++ show bufferData ++ " Got: " ++ show downloadedBufData
  
                                            if texMatch && bufMatch
                                              then
                                                let appRes = AppResources oTex tCopy tSmall oBuf bCopy originalBS bufferData imgW imgH
                                                 in return $ Just appRes
                                              else return Nothing -- Comparison failed
                                      else sdlLog "Wait for fence failed." >> return Nothing
              _ -> do
                -- Case for one of the initial Just oTex ... downTrans being Nothing
                sdlLog "!!! Failed to create one or more initial GPU resources."
                -- More robust cleanup would release any non-Nothing values from the tuple
                maybe (pure ()) (sdlReleaseGPUTexture contextDevice) maybeOrigTex
                maybe (pure ()) (sdlReleaseGPUTexture contextDevice) maybeTexCopy
                maybe (pure ()) (sdlReleaseGPUTexture contextDevice) maybeTexSmall
                maybe (pure ()) (sdlReleaseGPUBuffer contextDevice) maybeOrigBuf
                maybe (pure ()) (sdlReleaseGPUBuffer contextDevice) maybeBufCopy
                maybe (pure ()) (sdlReleaseGPUTransferBuffer contextDevice) maybeUploadTransfer
                maybe (pure ()) (sdlReleaseGPUTransferBuffer contextDevice) maybeDownloadTransfer
                return Nothing
  
  -- releaseResources
  releaseResources :: Context -> Maybe AppResources -> IO ()
  releaseResources _ Nothing = return ()
  releaseResources Context {..} (Just AppResources {..}) = do
    sdlLog "--> Releasing AppResources..."
    sdlReleaseGPUTexture contextDevice resOriginalTexture
    sdlReleaseGPUTexture contextDevice resTextureCopy
    sdlReleaseGPUTexture contextDevice resTextureSmall
    sdlReleaseGPUBuffer contextDevice resOriginalBuffer
    sdlReleaseGPUBuffer contextDevice resBufferCopy
    sdlLog "<-- AppResources Released."
  
  -- eventLoopGPU
  eventLoopGPU :: Context -> AppResources -> IO ()
  eventLoopGPU context resources = do
    sdlPumpEvents
    shouldQuitRef <- newIORef False
    processEventsGPU shouldQuitRef
    shouldQuit <- readIORef shouldQuitRef
    unless shouldQuit $ renderFrameGPU context resources >> eventLoopGPU context resources
  
  -- processEventsGPU & handleEventGPU (Simplified to just quit)
  processEventsGPU :: IORef Bool -> IO ()
  processEventsGPU sr = sdlPollEvent >>= maybe (pure ()) (handleEventGPU >=> \q -> when q (writeIORef sr True) >> processEventsGPU sr)
  
  handleEventGPU :: SDLEvent -> IO Bool
  handleEventGPU (SDLEventQuit _) = sdlLog "Quit." >> return True
  handleEventGPU (SDLEventKeyboard (SDLKeyboardEvent _ _ _ _ sc _ _ _ d _)) | d && sc == SDL_SCANCODE_Q = return True
  handleEventGPU _ = return False
  
  -- renderFrameGPU: Blits the processed textures to the screen
  renderFrameGPU :: Context -> AppResources -> IO ()
  renderFrameGPU Context {..} AppResources {..} = do
    maybeCmdbuf <- sdlAcquireGPUCommandBuffer contextDevice
    case maybeCmdbuf of
      Nothing -> sdlLog "Error: Failed to acquire render command buffer for drawing."
      Just cmdbuf -> do
        -- Use the helper function for dimensions
        maybeSwapResult <- sdlWaitAndAcquireGPUSwapchainTexture cmdbuf contextWindow
        case maybeSwapResult of
          Nothing -> void (sdlSubmitGPUCommandBuffer cmdbuf `finally` pure ())
          Just (swapchainTexture, swapW, swapH) -> do
            -- swapW, swapH are Word32
  
            -- Clear swapchain
            let colorTargetInfo =
                  defaultColorTargetInfo -- Use default
                    { texture = swapchainTexture, -- Update texture
                      loadOp = SDL_GPU_LOADOP_CLEAR,
                      storeOp = SDL_GPU_STOREOP_STORE,
                      clearColor = SDLFColor 0 0 0 1
                      -- targetCycle = False (already default)
                    }
            mClearPass <- sdlBeginGPURenderPass cmdbuf [colorTargetInfo] Nothing
            case mClearPass of
              Nothing -> sdlLog "Failed to begin clear pass"
              Just clearPass -> sdlEndGPURenderPass clearPass
  
            -- Blit OriginalTexture to top-left quadrant
            let blitInfoOrig =
                  SDLGPUBlitInfo
                    (defaultBlitRegion resOriginalTexture textureWidth textureHeight)
                    ((defaultBlitRegion swapchainTexture (fromIntegral swapW) (fromIntegral swapH)) {gpuBlitRegW = swapW `div` 2, gpuBlitRegH = swapH `div` 2})
                    SDL_GPU_LOADOP_LOAD
                    (SDLFColor 0 0 0 0)
                    SDL_FLIP_NONE
                    SDL_GPU_FILTER_NEAREST
                    False
            sdlBlitGPUTexture cmdbuf blitInfoOrig
  
            -- Blit TextureCopy to top-right quadrant
            let blitInfoCopy =
                  SDLGPUBlitInfo
                    (defaultBlitRegion resTextureCopy textureWidth textureHeight)
                    ((defaultBlitRegion swapchainTexture (fromIntegral swapW) (fromIntegral swapH)) {gpuBlitRegX = swapW `div` 2, gpuBlitRegW = swapW `div` 2, gpuBlitRegH = swapH `div` 2})
                    SDL_GPU_LOADOP_LOAD
                    (SDLFColor 0 0 0 0)
                    SDL_FLIP_NONE
                    SDL_GPU_FILTER_NEAREST
                    False
            sdlBlitGPUTexture cmdbuf blitInfoCopy
  
            -- Blit TextureSmall to bottom-center-ish quadrant
            let blitInfoSmall =
                  SDLGPUBlitInfo
                    (defaultBlitRegion resTextureSmall (textureWidth `div` 2) (textureHeight `div` 2))
                    ( (defaultBlitRegion swapchainTexture (fromIntegral swapW) (fromIntegral swapH))
                        { gpuBlitRegX = swapW `div` 4,
                          gpuBlitRegY = swapH `div` 2,
                          gpuBlitRegW = swapW `div` 2,
                          gpuBlitRegH = swapH `div` 2
                        }
                    )
                    SDL_GPU_LOADOP_LOAD
                    (SDLFColor 0 0 0 0)
                    SDL_FLIP_NONE
                    SDL_GPU_FILTER_NEAREST
                    False
            sdlBlitGPUTexture cmdbuf blitInfoSmall
  
            submitted <- sdlSubmitGPUCommandBuffer cmdbuf
            unless submitted $ sdlGetError >>= sdlLog . ("Draw Submit failed: " ++)
  