Bindings.hs

{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedRecordUpdate #-}

module Codec.HEIF.Bindings
  ( Image (..),
    Metadata (..),
    MetadataBlock (..),
    HeifException (..),
    HeifError (..),
    HeifSubError (..),
    MetadataType (..),
    Encoder (..),
    EncodeQuality (..),
    EncodeOptions (..),
    DecodeOptions (..),
    heifVersion,
    withContext,
    withPrimaryImage,
    withDecodedRGB,
    withRGBImage,
    withEncoder,
    withNewRGBPlane,
    withEncodeImage,
    encodeQuality,
    readByteStringIntoContext,
    readFileIntoContext,
    extractImageData,
    extractMetadata,
    fillPlaneFromByteString,
    addMetadataBlock,
    encodeToByteString,
    encodeWriteToFile,
  )
where

import Codec.HEIF.Bindings.Generated qualified as HEIF
import Codec.HEIF.Bindings.Generated.Safe qualified as HEIF
import Codec.HEIF.Error
import Codec.HEIF.Error.Internal
import Control.Exception
import Control.Monad (when)
import Data.ByteString qualified as BS
import Data.ByteString.Builder qualified as BSB
import Data.ByteString.Internal qualified as BSI
import Data.ByteString.Lazy qualified as BSL
import Data.ByteString.Unsafe qualified as BSU (unsafeUseAsCStringLen)
import Data.IORef
import Data.Map qualified as M
import Foreign qualified
import Foreign.C.ConstPtr (ConstPtr (..))
import Foreign.C.String qualified as Foreign (peekCString, withCString)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Ptr (plusPtr)
import HsBindgen.Runtime.LibC qualified as Bindgen
import HsBindgen.Runtime.Prelude qualified as Bindgen
import HsBindgen.Runtime.PtrConst qualified as Bindgen
import HsBindgen.Runtime.Support qualified as Bindgen
import Numeric.Natural
import Prelude

data DecodeOptions
  = DecodeOptions
  { -- | Do not apply metadata transformations such as crop or rotate.
    doIgnoreTransformations :: Bool,
    -- | Downsamples high color depth images to 8 bits.
    doConvertHDRTo8Bit :: Bool
  }

newtype EncodeOptions
  = EncodeOptions
  { eoQuality :: EncodeQuality
  }

constNullPtr :: Bindgen.PtrConst a
constNullPtr = Bindgen.unsafeFromPtr Foreign.nullPtr

-- | Returns the version of the underlying libheif C library.
heifVersion :: IO String
heifVersion = do
  c_str <- HEIF.heif_get_version
  Foreign.peekCString $ Bindgen.unsafeToPtr c_str

ccharPtrToString :: Bindgen.PtrConst Bindgen.CChar -> IO (Maybe String)
ccharPtrToString constPtr = do
  let rawPtr = Bindgen.unsafeToPtr constPtr
  if rawPtr == Foreign.nullPtr
    then pure Nothing
    else Just <$> Foreign.peekCString rawPtr

errorToException :: HEIF.Heif_error -> IO HeifException
errorToException HEIF.Heif_error {..} = do
  message <- ccharPtrToString heif_error_message
  pure $ LibraryException Library {..}
  where
    code = convertHeifErrorCode heif_error_code
    subcode = convertHeifSubErrorCode heif_error_subcode

throwIfError :: IO HEIF.Heif_error -> IO ()
throwIfError action = do
  err <- action
  when (HEIF.unwrapHeif_error_code (HEIF.heif_error_code err) /= 0) $
    errorToException err >>= throwIO

newtype Context = Context
  { _unContextHandle :: Foreign.Ptr HEIF.Heif_context
  }

newtype ImageHandle = ImageHandle (Bindgen.PtrConst HEIF.Heif_image_handle)

newtype WritableImageHandle = WritableImageHandle (Foreign.Ptr HEIF.Heif_image_handle)

newtype HeifImage = HeifImage (Bindgen.PtrConst HEIF.Heif_image)

newtype WritableImage = WritableImage (Foreign.Ptr HEIF.Heif_image)

newtype Encoder = Encoder (Foreign.Ptr HEIF.Heif_encoder)

data Plane = Plane
  { planePtr :: Foreign.Ptr Bindgen.Word8,
    planeStride :: Int,
    planeWidth :: Natural,
    planeHeight :: Natural
  }

-- | Lossy values must be in range 0-100.
data EncodeQuality = Lossless | Lossy Natural

withContext :: (Context -> IO a) -> IO a
withContext = bracket acquire release
  where
    acquire = Context <$> HEIF.heif_context_alloc
    release (Context ptr) = HEIF.heif_context_free ptr

withPrimaryImage :: Context -> (ImageHandle -> IO a) -> IO a
withPrimaryImage (Context contextPtr) = bracket acquire release
  where
    acquire = Foreign.alloca $ \imageHandleOutPtr -> do
      throwIfError $
        HEIF.heif_context_get_primary_image_handle
          contextPtr
          imageHandleOutPtr
      ImageHandle . Bindgen.unsafeFromPtr <$> Foreign.peek imageHandleOutPtr
    release (ImageHandle ptr) = HEIF.heif_image_handle_release ptr

withOptionsPointer :: (Foreign.Ptr HEIF.Heif_decoding_options -> IO a) -> IO a
withOptionsPointer =
  bracket acquire release
  where
    acquire = HEIF.heif_decoding_options_alloc
    release = HEIF.heif_decoding_options_free

withDecodedRGB :: DecodeOptions -> ImageHandle -> (HeifImage -> IO a) -> IO a
withDecodedRGB DecodeOptions {..} (ImageHandle handlePtr) = bracket acquire release
  where
    acquire = Foreign.alloca $ \imageOutPtr -> do
      withOptionsPointer $ \optionsPtr -> do
        Foreign.poke
          optionsPtr.heif_decoding_options_ignore_transformations
          (if doIgnoreTransformations then 1 else 0)
        Foreign.poke
          optionsPtr.heif_decoding_options_convert_hdr_to_8bit
          (if doConvertHDRTo8Bit then 1 else 0)
        throwIfError $
          HEIF.heif_decode_image
            handlePtr
            imageOutPtr
            HEIF.Heif_colorspace_RGB
            HEIF.Heif_chroma_interleaved_RGB
            (Bindgen.unsafeFromPtr optionsPtr)
      HeifImage . Bindgen.unsafeFromPtr <$> Foreign.peek imageOutPtr
    release (HeifImage ptr) = HEIF.heif_image_release ptr

withEncoder :: Context -> (Encoder -> IO a) -> IO a
withEncoder (Context contextPtr) = bracket acquire release
  where
    acquire = Foreign.alloca $ \encoderOutPtr -> do
      throwIfError $
        HEIF.heif_context_get_encoder_for_format
          contextPtr
          HEIF.Heif_compression_HEVC
          encoderOutPtr
      Encoder <$> Foreign.peek encoderOutPtr
    release (Encoder ptr) = HEIF.heif_encoder_release ptr

encodeQuality :: EncodeQuality -> Encoder -> IO ()
encodeQuality Lossless (Encoder encoderPtr) =
  throwIfError $ HEIF.heif_encoder_set_lossless encoderPtr 1
encodeQuality (Lossy quality) (Encoder encoderPtr) = do
  throwIfError $ HEIF.heif_encoder_set_lossless encoderPtr 0
  throwIfError $ HEIF.heif_encoder_set_lossy_quality encoderPtr (fromIntegral quality)

readByteStringIntoContext :: Context -> BS.ByteString -> IO ()
readByteStringIntoContext (Context contextPtr) bs =
  BSU.unsafeUseAsCStringLen bs $ \(cStrPtr, len) -> do
    let voidPtr = Bindgen.unsafeFromPtr (Foreign.castPtr cStrPtr)
        cSize = fromIntegral len :: Bindgen.CSize
    -- heif_context_read_from_memory_without_copy is not suitable because the underlying
    -- bytestring could be garbage collected, in which case we'd be feeding invalid data
    -- to the decoder...
    throwIfError $ HEIF.heif_context_read_from_memory contextPtr voidPtr cSize constNullPtr

readFileIntoContext :: Context -> FilePath -> IO ()
readFileIntoContext (Context contextPtr) filePath =
  Foreign.withCString filePath $ \filePathPtr ->
    throwIfError $
      HEIF.heif_context_read_from_file contextPtr (Bindgen.unsafeFromPtr filePathPtr) constNullPtr

-- | Raw encoding of RGB data.
data Image = Image
  { imagePixelsRGB :: BS.ByteString,
    imageWidth :: Natural,
    imageHeight :: Natural
  }
  deriving (Show, Eq)

extractImageData :: HeifImage -> IO Image
extractImageData (HeifImage imgPtr) = Foreign.alloca $ \strideOutPtr -> do
  rawPlanePtr <-
    HEIF.heif_image_get_plane_readonly imgPtr HEIF.Heif_channel_interleaved strideOutPtr

  -- TODO: we could check that
  -- Bindgen.unsafeToPtr rawPlanePtr /= Foreign.nullPtr
  strideCSize <- Foreign.peek strideOutPtr
  heightCInt <- HEIF.heif_image_get_height imgPtr HEIF.Heif_channel_interleaved
  widthCInt <- HEIF.heif_image_get_width imgPtr HEIF.Heif_channel_interleaved

  let stride = fromIntegral strideCSize :: Int
      imageHeight :: Natural = fromIntegral heightCInt
      imageWidth :: Natural = fromIntegral widthCInt
      rowBytes = imageWidth * 3 -- 3 bytes per pixel for RGB
      totalExpectedBytes = rowBytes * imageHeight
      sourcePtr = Foreign.castPtr (Bindgen.unsafeToPtr rawPlanePtr)

  -- Allocate exactly the memory needed for the final contiguous ByteString
  imagePixelsRGB <- BSI.create (fromIntegral totalExpectedBytes) $ \destinationPtr ->
    mapM_
      ( \row -> do
          let currentSrcRow = sourcePtr `plusPtr` (row * stride)
              currentDestRow = destinationPtr `plusPtr` (row * fromIntegral rowBytes)
          copyBytes currentDestRow currentSrcRow $ fromIntegral rowBytes
      )
      [0 .. fromIntegral imageHeight - 1]

  pure $ Image {..}

data MetadataType = Exif | Mime | OtherMetadataType String
  deriving (Show, Eq)

-- | Metadata of type e.g. Exif, mime, or anything else. The block contains raw metadata bytes,
-- it is the caller's responsability to validate and decode it. For Exif consider using
-- [hsexif](https://hackage.haskell.org/package/hsexif).
data MetadataBlock = MetadataBlock
  { metadataBlockType :: MetadataType,
    metadataBlockContentType :: String,
    metadataBlockData :: BS.ByteString
  }
  deriving (Show, Eq)

-- | Metadata found in the image.
newtype Metadata = Metadata
  { metadataBlocks :: M.Map Bindgen.Word32 MetadataBlock
  }
  deriving (Show, Eq)

getMetadataBlocks :: ImageHandle -> IO (M.Map Bindgen.Word32 MetadataBlock)
getMetadataBlocks imageHandle@(ImageHandle image) = do
  numBlocks <- HEIF.heif_image_handle_get_number_of_metadata_blocks image constNullPtr
  let count = fromIntegral numBlocks
  if count <= 0
    then pure mempty
    else do
      ids <- Foreign.allocaArray count $ \idsPtr -> do
        _ActualCount <-
          HEIF.heif_image_handle_get_list_of_metadata_block_IDs
            image
            constNullPtr
            idsPtr
            (fromIntegral count)
        Foreign.peekArray count idsPtr
      M.fromList <$> mapM (extractMetadataBlock imageHandle) ids

-- | Exif metadata has a 4 0-byte prefix, strip it
dropExifHeader :: String -> BS.ByteString -> BS.ByteString
dropExifHeader "Exif" rawData
  | BS.length rawData > 4 = BS.drop 4 rawData
  | otherwise = rawData
dropExifHeader "mime" rawData = rawData -- Direct raw XML, no offset to drop
dropExifHeader _ rawData = rawData

mkMetadataType :: String -> MetadataType
mkMetadataType "Exif" = Exif
mkMetadataType "mime" = Mime
mkMetadataType other = OtherMetadataType other

metadataTypeToString :: MetadataType -> String
metadataTypeToString Exif = "Exif"
metadataTypeToString Mime = "mime"
metadataTypeToString (OtherMetadataType other) = other

extractMetadataBlock :: ImageHandle -> HEIF.Heif_item_id -> IO (Bindgen.Word32, MetadataBlock)
extractMetadataBlock (ImageHandle image) itemId = do
  -- Fetch the block type string ("Exif", "mime", etc.)
  metadataTypePtr <- HEIF.heif_image_handle_get_metadata_type image itemId
  metadataTypeString <- Foreign.peekCString (Bindgen.unsafeToPtr metadataTypePtr)

  metadataContentTypePtr <- HEIF.heif_image_handle_get_metadata_content_type image itemId
  metadataContentTypeString <- Foreign.peekCString (Bindgen.unsafeToPtr metadataContentTypePtr)

  -- Get exact size in bytes needed for memory allocation
  dataSize <- HEIF.heif_image_handle_get_metadata_size image itemId

  -- Allocate byte buffer, call C function, and pack into a ByteString
  bsData <- Foreign.allocaBytes (fromIntegral dataSize) $ \outBuf -> do
    throwIfError $ HEIF.heif_image_handle_get_metadata image itemId (Foreign.castPtr outBuf)
    -- Check 'err' here if your bindings expose error validation
    BS.packCStringLen (Foreign.castPtr outBuf, fromIntegral dataSize)

  pure
    ( HEIF.unwrapHeif_item_id itemId,
      MetadataBlock
        { metadataBlockType = mkMetadataType metadataTypeString,
          metadataBlockData = dropExifHeader metadataTypeString bsData,
          metadataBlockContentType = metadataContentTypeString
        }
    )

-- | Returns (some) metadata information about the image.
extractMetadata :: ImageHandle -> IO Metadata
extractMetadata imageHandle = do
  metadataBlocks <- getMetadataBlocks imageHandle
  pure $ Metadata {..}

withRGBImage :: Natural -> Natural -> (WritableImage -> IO a) -> IO a
withRGBImage width height = bracket acquire release
  where
    acquire = Foreign.alloca $ \imageOutPtr -> do
      throwIfError $
        HEIF.heif_image_create
          (fromIntegral width)
          (fromIntegral height)
          HEIF.Heif_colorspace_RGB
          HEIF.Heif_chroma_interleaved_RGB
          imageOutPtr
      WritableImage <$> Foreign.peek imageOutPtr
    release (WritableImage imageOutPtr) =
      HEIF.heif_image_release (Bindgen.unsafeFromPtr imageOutPtr)

withNewRGBPlane :: WritableImage -> Natural -> Natural -> (Plane -> IO a) -> IO a
withNewRGBPlane (WritableImage imagePtr) planeWidth planeHeight = bracket acquire release
  where
    acquire = Foreign.alloca $ \stridePtr -> do
      throwIfError $
        HEIF.heif_image_add_plane
          imagePtr
          HEIF.Heif_channel_interleaved
          (fromIntegral planeWidth)
          (fromIntegral planeHeight)
          24 -- 24 bits per pixel (RGB)
      planePtr <- HEIF.heif_image_get_plane imagePtr HEIF.Heif_channel_interleaved stridePtr
      planeStride <- fromIntegral <$> Foreign.peek stridePtr
      pure Plane {..}
    release _ = pure ()

fillPlaneFromByteString :: Plane -> BS.ByteString -> IO ()
fillPlaneFromByteString Plane {..} rgbData
  | dataLength /= expectedSize =
      throw $
        DecodeSizeMismatch
          SizeMismatch
            { expectedBytes = fromIntegral expectedSize,
              infoWidth = planeWidth,
              infoHeight = planeHeight,
              gotBytes = fromIntegral dataLength
            }
  | planeStride == rgbWidth = copyFast
  | otherwise = copyStrided
  where
    w = fromIntegral planeWidth :: Int
    h = fromIntegral planeHeight :: Int
    rgbWidth = w * 3
    expectedSize = rgbWidth * h
    dataLength = BS.length rgbData
    copyFast = BSU.unsafeUseAsCStringLen rgbData $ \(sourcePtr, len) ->
      copyBytes planePtr (Foreign.castPtr sourcePtr) len
    copyStrided = BSU.unsafeUseAsCStringLen rgbData $ \(srcPtr, _) -> do
      mapM_ (copyRow srcPtr) [0 .. h - 1]
    copyRow srcBytePtr y = do
      let srcRow = srcBytePtr `plusPtr` (y * rgbWidth)
          dstRow = planePtr `plusPtr` (y * planeStride)
      copyBytes dstRow srcRow rgbWidth

withEncodeImage ::
  Context -> WritableImage -> Encoder -> (WritableImageHandle -> IO a) -> IO a
withEncodeImage (Context contextPtr) (WritableImage image) (Encoder encoderPtr) =
  bracket acquire release
  where
    acquire = Foreign.alloca $ \imageHandlePtr -> do
      throwIfError $
        HEIF.heif_context_encode_image
          contextPtr
          (Bindgen.unsafeFromPtr image)
          encoderPtr
          constNullPtr
          imageHandlePtr
      WritableImageHandle <$> Foreign.peek imageHandlePtr

    release (WritableImageHandle imageHandlePtr) =
      HEIF.heif_image_handle_release (Bindgen.unsafeFromPtr imageHandlePtr)

-- | It's complementary to 'dropExifHeader'.
addExifHeader :: MetadataType -> BS.ByteString -> BS.ByteString
addExifHeader Exif = (BS.replicate 4 0x00 <>)
addExifHeader _ = id

addMetadataBlock :: Context -> WritableImageHandle -> MetadataBlock -> IO ()
addMetadataBlock (Context contextPtr) (WritableImageHandle imageHandle) MetadataBlock {..} =
  BSU.unsafeUseAsCStringLen (addExifHeader metadataBlockType metadataBlockData) $
    \(dataPtr, len) -> do
      Foreign.withCString (metadataTypeToString metadataBlockType) $ \metadataTypePtr ->
        Foreign.withCString metadataBlockContentType $ \contentTypePtr ->
          throwIfError $
            HEIF.heif_context_add_generic_metadata
              contextPtr
              (Bindgen.unsafeFromPtr imageHandle)
              (Bindgen.unsafeFromPtr $ Foreign.castPtr dataPtr)
              (fromIntegral len)
              (Bindgen.unsafeFromPtr metadataTypePtr)
              (Bindgen.unsafeFromPtr contentTypePtr)

type HsWriteCb =
  Bindgen.Ptr HEIF.Heif_context ->
  Bindgen.PtrConst Bindgen.Void ->
  Bindgen.CSize ->
  Bindgen.Ptr Bindgen.Void ->
  IO Bindgen.CInt

foreign import ccall "get_c_trampoline_ptr"
  c_get_trampoline_ptr :: IO (Bindgen.Ptr ())

foreign import ccall "wrapper"
  mkHsWriteCb :: HsWriteCb -> IO (Bindgen.FunPtr HsWriteCb)

encodeWriteToFile :: Context -> FilePath -> IO ()
encodeWriteToFile (Context contextPtr) filePath =
  Foreign.withCString filePath $ \filePathPtr -> do
    throwIfError $
      HEIF.heif_context_write_to_file
        contextPtr
        (Bindgen.unsafeFromPtr filePathPtr)

encodeToByteString :: Context -> IO BSL.ByteString
encodeToByteString (Context contextPtr) = do
  bufferRef <- newIORef mempty
  bracket (acquire bufferRef) release (action bufferRef)
  where
    acquire ref = mkHsWriteCb (writeCallback ref)
    release = Foreign.freeHaskellFunPtr
    action ref chunkWriteFunctionPtr = do
      Foreign.alloca $ \writerPtr -> do
        trampolinePtr <- c_get_trampoline_ptr
        Foreign.poke writerPtr.heif_writer_writer_api_version (1 :: Bindgen.CInt)
        Foreign.poke writerPtr.heif_writer_write (Foreign.castPtrToFunPtr trampolinePtr)

        -- Trigger the encoding; this will invoke `cbPtr` multiple times
        -- Critical: pass the Haskell FunPtr as the 'userdata' argument!
        throwIfError $
          HEIF.heif_context_write
            contextPtr
            writerPtr
            (Foreign.castFunPtrToPtr chunkWriteFunctionPtr)

      -- Extract and finalize the ByteString from the builder
      finalBuilder <- readIORef ref
      pure $ BSB.toLazyByteString finalBuilder

    -- The Callback: appends raw C memory chunks to our IORef Builder
    writeCallback :: IORef BSB.Builder -> HsWriteCb
    writeCallback ref _ctx dataPtr size _userData = do
      chunk <- BS.packCStringLen (Foreign.castPtr $ Bindgen.unsafeToPtr dataPtr, fromIntegral size)
      modifyIORef' ref (\b -> b <> BSB.byteString chunk)
      pure 0