content: {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- |
  -- Example     : GPUTonemappingExample
  -- Description : Demonstrates HDR image loading, compute shader-based tonemapping, and color space conversion.
  -- Copyright   : (c) Kyle Lukaszek, 2025
  -- License     : BSD3
  --
  -- Based on the SDL_gpu_examples/Tonemapping C example.
  -- This example showcases a comprehensive GPU pipeline:
  -- - Loading an HDR (High Dynamic Range) image file (e.g., .hdr) using a wrapper around stb_image's `stbi_loadf`.
  -- - Creating multiple textures for different stages:
  --     - `HDRTexture`: Stores the original HDR image data (e.g., R32G32B32A32_FLOAT format).
  --     - `ToneMapTexture`: Intermediate texture to store the tonemapped HDR data (e.g., R16G16B16A16_FLOAT).
  --     - `TransferTexture`: Intermediate texture for color space converted output, matching swapchain format.
  -- - Uploading the loaded HDR pixel data to the `HDRTexture`.
  -- - Creating several compute pipelines:
  --     - A set of tonemapping operators (Reinhard, ACES, etc.), each implemented as a compute shader. These shaders
  --       read from `HDRTexture` and write the tonemapped result to `ToneMapTexture`.
  --     - `LinearToSRGBPipeline`: A compute shader to convert linear HDR data (from `ToneMapTexture`) to sRGB space,
  --       writing to `TransferTexture`.
  --     - `LinearToST2084Pipeline`: A compute shader to convert linear HDR data (from `ToneMapTexture`) to ST2084 (PQ)
  --       HDR10 space, writing to `TransferTexture`.
  -- - Managing user interaction to cycle through:
  --     - Different SDL Swapchain Compositions (SDR, SDR_LINEAR, HDR_EXTENDED_LINEAR, HDR10_ST2084) using
  --       `SDL_SetGPUSwapchainParameters` and checking support with `SDL_WindowSupportsGPUSwapchainComposition`.
  --     - Different tonemapping operators.
  -- - In each frame's render process:
  --     1. **Tonemapping Pass:**
  --         - A compute pass is initiated.
  --         - The selected tonemapping compute pipeline is bound.
  --         - `HDRTexture` is bound as a read-only storage texture input.
  --         - `ToneMapTexture` is bound as a read/write storage texture output.
  --         - The compute shader is dispatched.
  --     2. **Color Space Conversion Pass (Conditional):**
  --         - If the current swapchain composition is SDR or HDR10_ST2084, another compute pass is run.
  --         - The appropriate conversion pipeline (`LinearToSRGBPipeline` or `LinearToST2084Pipeline`) is bound.
  --         - `ToneMapTexture` is bound as a read-only storage texture input.
  --         - `TransferTexture` is bound as a read/write storage texture output.
  --         - The compute shader is dispatched.
  --     3. **Blit to Swapchain:**
  --         - The resulting texture (`ToneMapTexture` or `TransferTexture` if conversion occurred) is blitted to the
  --           acquired swapchain texture using `SDL_BlitGPUTexture`.
  -- - This example demonstrates a multi-stage image processing pipeline using compute shaders, handling
  --   different color spaces, and adapting output based on swapchain capabilities.
  -- |
  module Main where
  
  import Control.Exception (bracket, bracketOnError, finally)
  import Control.Monad (forM, forM_, unless, void, when)
  import Data.Bits ((.|.))
  import Data.IORef
  import Data.Maybe (catMaybes, fromJust, isJust, isNothing)
  import Foreign.Ptr (castPtr)
  import GPUCommon
  import SDL
  import System.Exit (exitFailure, exitSuccess)
  
  -- AppResources
  data AppResources = AppResources
    { resHDRTexture :: SDLGPUTexture,
      resToneMapTexture :: SDLGPUTexture,
      resTransferTexture :: SDLGPUTexture,
      resTonemapOperators :: [SDLGPUComputePipeline], -- Store all loaded operators
      resLinearToSRGBPipeline :: SDLGPUComputePipeline,
      resLinearToST2084Pipeline :: SDLGPUComputePipeline,
      -- User interaction state
      currentSwapchainCompositionRef :: IORef SDLGPUSwapchainComposition,
      currentTonemapOperatorRef :: IORef SDLGPUComputePipeline,
      swapchainCompositionIndexRef :: IORef Int,
      tonemapOperatorIndexRef :: IORef Int
    }
  
  -- Global constants for UI cycling (matching C example)
  swapchainCompositions :: [SDLGPUSwapchainComposition]
  swapchainCompositions =
    [ SDL_GPU_SWAPCHAINCOMPOSITION_SDR,
      SDL_GPU_SWAPCHAINCOMPOSITION_SDR_LINEAR,
      SDL_GPU_SWAPCHAINCOMPOSITION_HDR_EXTENDED_LINEAR,
      SDL_GPU_SWAPCHAINCOMPOSITION_HDR10_ST2084
    ]
  
  swapchainCompositionNames :: [String]
  swapchainCompositionNames =
    [ "SDR",
      "SDR_LINEAR",
      "HDR_EXTENDED_LINEAR",
      "HDR10_ST2084"
    ]
  
  tonemapOperatorShaderFiles :: [String] -- Base filenames for compute shaders
  tonemapOperatorShaderFiles =
    [ "ToneMapReinhard.comp",
      "ToneMapExtendedReinhardLuminance.comp",
      "ToneMapHable.comp",
      "ToneMapACES.comp"
    ]
  
  tonemapOperatorDisplayNames :: [String]
  tonemapOperatorDisplayNames = ["Reinhard", "ExtendedReinhardLuminance", "Hable", "ACES"]
  
  -- main
  main :: IO ()
  main = do
    sdlLog $ "Compiled SDL Version: " ++ show sdlVersion
    linkedVersion <- sdlGetVersion
    sdlLog $ "Linked SDL Version: " ++ show linkedVersion
  
    maybeResult <- withContext "SDL3 Haskell GPU Tonemapping" [] 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
      (createResources context)
      (releaseResources context)
      $ \case
        Nothing -> sdlLog "Failed to create resources. Exiting."
        Just resources -> do
          sdlLog "Resources created successfully."
          sdlLog "Press Left/Right to cycle swapchain composition"
          sdlLog "Press Up/Down to cycle tonemap operators"
          eventLoopGPU context resources
  
  -- createResources
  createResources :: Context -> IO (Maybe AppResources)
  createResources Context {..} = do
    sdlLog "--- Beginning Resource Creation ---"
    let desiredChannels = 4
  
    -- Load HDR Image Data
    (w, h) <-
      bracketOnError
        (loadHDRImageFromFile "memorial.hdr" desiredChannels)
        ( \case
            Just (ptr, _, _, _) -> free (castPtr ptr)
            _ -> pure ()
        )
        ( \case
            Nothing -> sdlLog "Could not load HDR image data!" >> error "HDR Load Failed" -- Force exit
            Just (_, w, h, n) -> do
              sdlLog $ "HDR Image loaded: " ++ show w ++ "x" ++ show h ++ ", channels: " ++ show n
              -- Adjust window size based on loaded HDR image
              _ <- sdlSetWindowSize contextWindow w h
              return (w, h)
        )
  
    -- Create Textures
    let hdrFormat = SDL_GPU_TEXTUREFORMAT_R32G32B32A32_FLOAT
    let toneMapFormat = SDL_GPU_TEXTUREFORMAT_R16G16B16A16_FLOAT
    swapchainFmt <- sdlGetGPUSwapchainTextureFormat contextDevice contextWindow
  
    maybeHDRTexture <-
      sdlCreateGPUTexture
        contextDevice
        ( (defaultTextureCreateInfo w h)
            { texInfoFormat = hdrFormat,
              texInfoUsage = SDL_GPU_TEXTUREUSAGE_SAMPLER .|. SDL_GPU_TEXTUREUSAGE_COMPUTE_STORAGE_READ -- Read by tonemapper
            }
        )
    maybeToneMapTexture <-
      sdlCreateGPUTexture
        contextDevice
        ( (defaultTextureCreateInfo w h)
            { texInfoFormat = toneMapFormat,
              texInfoUsage = SDL_GPU_TEXTUREUSAGE_SAMPLER .|. SDL_GPU_TEXTUREUSAGE_COMPUTE_STORAGE_READ .|. SDL_GPU_TEXTUREUSAGE_COMPUTE_STORAGE_WRITE -- Read by post, written by tonemapper
            }
        )
    maybeTransferTexture <-
      sdlCreateGPUTexture
        contextDevice
        ( (defaultTextureCreateInfo w h)
            { texInfoFormat = swapchainFmt,
              texInfoUsage = SDL_GPU_TEXTUREUSAGE_SAMPLER .|. SDL_GPU_TEXTUREUSAGE_COMPUTE_STORAGE_WRITE -- Written by post, read by blit
            }
        )
  
    -- Create Compute Pipelines for Tonemapping
    let tonemapBaseCI =
          defaultComputePipelineCreateInfo
            { numReadOnlyStorageTextures = 1, -- Reads HDRTexture
              numReadWriteStorageTextures = 1, -- Writes ToneMapTexture
              threadCountX = 8,
              threadCountY = 8,
              threadCountZ = 1
            }
    -- Use forM to attempt creation of all tonemap operators
    tonemapOps <- forM tonemapOperatorShaderFiles $ \shaderFile ->
      createComputePipelineFromShader contextDevice shaderFile tonemapBaseCI
  
    -- Create Post-Processing Compute Pipelines
    let postProcessBaseCI =
          defaultComputePipelineCreateInfo
            { numReadOnlyStorageTextures = 1, -- Reads ToneMapTexture
              numReadWriteStorageTextures = 1, -- Writes TransferTexture
              threadCountX = 8,
              threadCountY = 8,
              threadCountZ = 1
            }
    maybeLinearToSRGB <- createComputePipelineFromShader contextDevice "LinearToSRGB.comp" postProcessBaseCI
    maybeLinearToST2084 <- createComputePipelineFromShader contextDevice "LinearToST2084.comp" postProcessBaseCI
  
    -- Upload HDR Image Data to HDRTexture
    -- (Re-load for upload, or pass ptr around; re-loading is simpler for bracket scope)
    uploadSuccess <-
      bracketOnError
        (loadHDRImageFromFile "memorial.hdr" 4)
        ( \case
            Just (ptr, _, _, _) -> free (castPtr ptr)
            _ -> pure ()
        )
        ( \case
            Nothing -> sdlLog "Failed to re-load HDR for upload." >> return False
            Just (pixelPtr, _, _, _) ->
              case maybeHDRTexture of
                Nothing -> sdlLog "HDR Texture not created, cannot upload." >> return False
                Just hdrTex -> do
                  let dataSizeBytes = fromIntegral (w * h * desiredChannels * 4) -- 4 bytes per pixel
                  bracket
                    (createTransferBuffer contextDevice dataSizeBytes SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD "HDRUpload")
                    (cleanupTransferBuffer contextDevice)
                    $ \case
                      Nothing -> return False
                      Just transferBuf -> do
                        mapOk <- bracket
                          (sdlMapGPUTransferBuffer contextDevice transferBuf False)
                          (\mptr -> when (isJust mptr) $ sdlUnmapGPUTransferBuffer contextDevice transferBuf)
                          $ \case
                            Nothing -> return False
                            Just mappedPtr -> memcpy mappedPtr (castPtr pixelPtr) (fromIntegral dataSizeBytes) >> return True
                        if mapOk
                          then do
                            cmdBuf <- sdlAcquireGPUCommandBuffer contextDevice
                            case cmdBuf of
                              Nothing -> return False
                              Just cb -> do
                                mcp <- sdlBeginGPUCopyPass cb
                                case mcp of
                                  Nothing -> return False
                                  Just cp -> do
                                    let texSrc = SDLGPUTextureTransferInfo transferBuf 0 (fromIntegral w) (fromIntegral h)
                                    let texDst = defaultTextureRegion hdrTex w h
                                    sdlUploadToGPUTexture cp texSrc texDst False
                                    sdlEndGPUCopyPass cp
                                    sdlSubmitGPUCommandBuffer cb >>= \s -> if s then sdlWaitForGPUIdle contextDevice >> return True else return False
                          else return False
        )
  
    -- Check all creations
    if not uploadSuccess || any isNothing tonemapOps || isNothing maybeLinearToSRGB || isNothing maybeLinearToST2084 || isNothing maybeHDRTexture || isNothing maybeToneMapTexture || isNothing maybeTransferTexture
      then do
        sdlLog "!!! Failed to create one or more critical resources."
        -- Manual cleanup of successfully created resources before this point
        maybe (pure ()) (sdlReleaseGPUTexture contextDevice) maybeHDRTexture
        maybe (pure ()) (sdlReleaseGPUTexture contextDevice) maybeToneMapTexture
        maybe (pure ()) (sdlReleaseGPUTexture contextDevice) maybeTransferTexture
        forM_ (catMaybes tonemapOps) (sdlReleaseGPUComputePipeline contextDevice)
        maybe (pure ()) (sdlReleaseGPUComputePipeline contextDevice) maybeLinearToSRGB
        maybe (pure ()) (sdlReleaseGPUComputePipeline contextDevice) maybeLinearToST2084
        return Nothing
      else do
        let allTonemapOps = catMaybes tonemapOps
        -- Initialize IORefs for UI state
        swapCompIdxRef <- newIORef 0
        toneOpIdxRef <- newIORef 0
        currentSwapCompRef <- newIORef (head swapchainCompositions)
        currentToneOpRef <- newIORef (head allTonemapOps)
  
        return $
          Just
            AppResources
              { resHDRTexture = fromJust maybeHDRTexture,
                resToneMapTexture = fromJust maybeToneMapTexture,
                resTransferTexture = fromJust maybeTransferTexture,
                resTonemapOperators = allTonemapOps,
                resLinearToSRGBPipeline = fromJust maybeLinearToSRGB,
                resLinearToST2084Pipeline = fromJust maybeLinearToST2084,
                currentSwapchainCompositionRef = currentSwapCompRef,
                currentTonemapOperatorRef = currentToneOpRef,
                swapchainCompositionIndexRef = swapCompIdxRef,
                tonemapOperatorIndexRef = toneOpIdxRef
              }
  
  -- releaseResources
  releaseResources :: Context -> Maybe AppResources -> IO ()
  releaseResources _ Nothing = return ()
  releaseResources Context {..} appResourcesMaybe =
    -- Use pattern match on Maybe
    case appResourcesMaybe of
      Just AppResources {..} -> do
        let dev = contextDevice -- For convenience
        sdlLog "--> Releasing AppResources..."
  
        sdlLog $ "  --> Releasing " ++ show (length resTonemapOperators) ++ " Tonemap Operator Compute Pipelines..."
        forM_ (zip [0 ..] resTonemapOperators) $ \(idx, pipeline) -> do
          -- It might be nice to associate names with these if you store them that way
          sdlLog $ "    Releasing Tonemap Operator Pipeline #" ++ show idx ++ ": " ++ show (tonemapOperatorDisplayNames !! idx) ++ " " ++ show pipeline
          sdlReleaseGPUComputePipeline dev pipeline
  
        sdlLog $ "  --> Releasing LinearToSRGB Compute Pipeline: " ++ show resLinearToSRGBPipeline
        sdlReleaseGPUComputePipeline dev resLinearToSRGBPipeline
  
        sdlLog $ "  --> Releasing LinearToST2084 Compute Pipeline: " ++ show resLinearToST2084Pipeline
        sdlReleaseGPUComputePipeline dev resLinearToST2084Pipeline
  
        sdlLog $ "  --> Releasing HDR Texture: " ++ show resHDRTexture
        sdlReleaseGPUTexture dev resHDRTexture
  
        sdlLog $ "  --> Releasing ToneMap Texture: " ++ show resToneMapTexture
        sdlReleaseGPUTexture dev resToneMapTexture
  
        sdlLog $ "  --> Releasing Transfer Texture: " ++ show resTransferTexture
        sdlReleaseGPUTexture dev resTransferTexture
  
        sdlLog "<-- AppResources Released."
  
  -- eventLoopGPU
  eventLoopGPU :: Context -> AppResources -> IO ()
  eventLoopGPU context resources = do
    sdlPumpEvents
    shouldQuitRef <- newIORef False
    processEventsGPU context resources shouldQuitRef
    shouldQuit <- readIORef shouldQuitRef
    unless shouldQuit $ do
      renderFrameGPU context resources
      eventLoopGPU context resources
  
  -- processEventsGPU
  processEventsGPU :: Context -> AppResources -> IORef Bool -> IO ()
  processEventsGPU context resources quitRef = do
    maybeEvent <- sdlPollEvent
    case maybeEvent of
      Nothing -> return ()
      Just event -> do
        quit <- handleEventGPU context resources event
        when quit $ writeIORef quitRef True
        processEventsGPU context resources quitRef
  
  -- handleEventGPU for UI interaction
  handleEventGPU :: Context -> AppResources -> SDLEvent -> IO Bool
  handleEventGPU context@Context {} AppResources {..} event = case event of
    SDLEventQuit _ -> sdlLog "Quit event received." >> return True
    SDLEventKeyboard (SDLKeyboardEvent _ _ _ _ scancode _ _ _ down _) | down -> do
      case scancode of
        SDL_SCANCODE_Q -> return True
        SDL_SCANCODE_LEFT -> do
          idx <- atomicModifyIORef' swapchainCompositionIndexRef $ \i ->
            let newIdx = (i - 1 + length swapchainCompositions) `mod` length swapchainCompositions
             in (newIdx, newIdx)
          changeSwapchainComposition context (swapchainCompositions !! idx) (swapchainCompositionNames !! idx) currentSwapchainCompositionRef
          return False
        SDL_SCANCODE_RIGHT -> do
          idx <- atomicModifyIORef' swapchainCompositionIndexRef $ \i ->
            let newIdx = (i + 1) `mod` length swapchainCompositions
             in (newIdx, newIdx)
          changeSwapchainComposition context (swapchainCompositions !! idx) (swapchainCompositionNames !! idx) currentSwapchainCompositionRef
          return False
        SDL_SCANCODE_UP -> do
          idx <- atomicModifyIORef' tonemapOperatorIndexRef $ \i ->
            let newIdx = (i - 1 + length resTonemapOperators) `mod` length resTonemapOperators
             in (newIdx, newIdx)
          changeTonemapOperator (resTonemapOperators !! idx) (tonemapOperatorDisplayNames !! idx) currentTonemapOperatorRef
          return False
        SDL_SCANCODE_DOWN -> do
          idx <- atomicModifyIORef' tonemapOperatorIndexRef $ \i ->
            let newIdx = (i + 1) `mod` length resTonemapOperators
             in (newIdx, newIdx)
          changeTonemapOperator (resTonemapOperators !! idx) (tonemapOperatorDisplayNames !! idx) currentTonemapOperatorRef
          return False
        _ -> return False
    _ -> return False
  
  -- UI Change Helpers
  changeSwapchainComposition :: Context -> SDLGPUSwapchainComposition -> String -> IORef SDLGPUSwapchainComposition -> IO ()
  changeSwapchainComposition Context {..} newComp compName compRef = do
    supported <- sdlWindowSupportsGPUSwapchainComposition contextDevice contextWindow newComp
    if supported
      then do
        sdlLog $ "Changing swapchain composition to " ++ compName
        success <- sdlSetGPUSwapchainParameters contextDevice contextWindow newComp SDL_GPU_PRESENTMODE_VSYNC
        when success $ writeIORef compRef newComp
        unless success $ sdlGetError >>= sdlLog . ("Failed to set swapchain params: " ++)
      else
        sdlLog $ "Swapchain composition " ++ compName ++ " unsupported."
  
  changeTonemapOperator :: SDLGPUComputePipeline -> String -> IORef SDLGPUComputePipeline -> IO ()
  changeTonemapOperator newOp opName opRef = do
    sdlLog $ "Changing tonemap operator to " ++ opName
    writeIORef opRef newOp
  
  -- renderFrameGPU
  renderFrameGPU :: Context -> AppResources -> IO ()
  renderFrameGPU Context {..} AppResources {..} = do
    maybeCmdbuf <- sdlAcquireGPUCommandBuffer contextDevice
    case maybeCmdbuf of
      Nothing -> sdlLog "Error: Failed to acquire render command buffer."
      Just cmdbuf -> do
        maybeSwapResult <- sdlWaitAndAcquireGPUSwapchainTexture cmdbuf contextWindow
        case maybeSwapResult of
          Nothing -> void (sdlSubmitGPUCommandBuffer cmdbuf `finally` pure ())
          Just (swapchainTexture, swapW, swapH) -> do
            let texW = swapW
            let texH = swapH
            -- 1. Tonemap Pass
            currentToneOp <- readIORef currentTonemapOperatorRef
            maybeTonePass <-
              sdlBeginGPUComputePass
                cmdbuf
                [SDLGPUStorageTextureReadWriteBinding resToneMapTexture 0 0 True] -- Write to ToneMapTexture, cycle
                [] -- No storage buffers
            case maybeTonePass of
              Nothing -> sdlLog "!!! Failed to begin Tonemap compute pass."
              Just tonePass -> do
                sdlBindGPUComputePipeline tonePass currentToneOp
                -- Bind HDRTexture as read-only for the tonemapper
                -- Assuming slot 0 for read-only storage textures in tonemapper shader
                sdlBindGPUComputeStorageTextures tonePass 0 [resHDRTexture]
                sdlDispatchGPUCompute tonePass (texW `div` 8) (texH `div` 8) 1
                sdlEndGPUComputePass tonePass
  
            -- 2. Transfer to Target Color Space if Necessary
            currentSwapComp <- readIORef currentSwapchainCompositionRef
            let blitSourceTexture =
                  if currentSwapComp == SDL_GPU_SWAPCHAINCOMPOSITION_SDR
                    || currentSwapComp == SDL_GPU_SWAPCHAINCOMPOSITION_HDR10_ST2084
                    then do
                      -- This block needs to be IO (Maybe SDLGPUTexture)
                      mp <-
                        sdlBeginGPUComputePass
                          cmdbuf
                          [SDLGPUStorageTextureReadWriteBinding resTransferTexture 0 0 True]
                          []
                      case mp of
                        Nothing -> sdlLog "!!! Failed to begin Color Space transform pass." >> return resToneMapTexture -- Fallback
                        Just postPass -> do
                          let pipelineToUse =
                                if currentSwapComp == SDL_GPU_SWAPCHAINCOMPOSITION_SDR
                                  then resLinearToSRGBPipeline
                                  else resLinearToST2084Pipeline
                          sdlBindGPUComputePipeline postPass pipelineToUse
                          -- Bind ToneMapTexture as read-only for this pass
                          sdlBindGPUComputeStorageTextures postPass 0 [resToneMapTexture]
                          sdlDispatchGPUCompute postPass (texW `div` 8) (texH `div` 8) 1
                          sdlEndGPUComputePass postPass
                          return resTransferTexture -- Success, blit from TransferTexture
                    else return resToneMapTexture -- No transform needed, blit from ToneMapTexture
            finalBlitSource <- blitSourceTexture -- Execute the IO action
  
            -- 3. Blit to Swapchain
            let srcBlitRegion = defaultBlitRegion finalBlitSource (fromIntegral texW) (fromIntegral texH)
            let dstBlitRegion =
                  (defaultBlitRegion swapchainTexture (fromIntegral swapW) (fromIntegral swapH))
                    { gpuBlitRegW = swapW,
                      gpuBlitRegH = swapH
                    }
  
            let blitInfo =
                  SDLGPUBlitInfo
                    { gpuBlitInfoSource = srcBlitRegion,
                      gpuBlitInfoDestination = dstBlitRegion,
                      gpuBlitInfoLoadOp = SDL_GPU_LOADOP_DONT_CARE,
                      gpuBlitInfoClearColor = SDLFColor 0 0 0 0,
                      gpuBlitInfoFlipMode = SDL_FLIP_NONE,
                      gpuBlitInfoFilter = SDL_GPU_FILTER_NEAREST, -- C example uses NEAREST for final blit
                      gpuBlitInfoCycle = False
                    }
            sdlBlitGPUTexture cmdbuf blitInfo
  
            -- Submit
            submitted <- sdlSubmitGPUCommandBuffer cmdbuf
            unless submitted $ sdlGetError >>= sdlLog . ("Submit failed: " ++)
  