{-# LANGUAGE ForeignFunctionInterface, NamedFieldPuns #-}

module X11.Image (
  PlaneMask,
  allPlanes,
  Image(..),
  getImage,
  getPixel,
  writeBitmapFile,
  getHeader,
  readColors,
  getXwdColor,
  freeReadColors,
  saveImage
) where

import Control.Exception
import Control.Monad
import Foreign
import Foreign.C.String
import Foreign.C.Types
import GHC.IO.Exception
import Graphics.X11.Xlib hiding (Image)
import Graphics.X11.Xlib.Extras
import System.IO

#include <X11/Xlib.h>
#include <X11/XWDFile.h>

type PlaneMask = CULong

foreign import ccall unsafe "XAllPlanes" allPlanes :: IO PlaneMask

data Image = Image {
  width :: !Dimension,
  height :: !Dimension,
  xOffset :: !Position,
  format :: !ImageFormat,
  imageData :: !(Ptr CChar),
  byteOrder :: !ByteOrder,
  imageBitmapUnit :: !CInt,
  imageBitmapBitOrder :: !ByteOrder,
  imageBitmapPad :: !CInt,
  depth :: !CInt,
  bytesPerLine :: !CInt,
  bitsPerPixel :: !CInt,
  redMask :: !PlaneMask,
  greenMask :: !PlaneMask,
  blueMask :: !PlaneMask
}

#let alignment t = "%lu", (unsigned long)offsetof(struct {char x__; t (y__); }, y__)

instance Storable Image where
  sizeOf _ = #{size XImage}
  alignment _ = #{alignment XImage}
  peek ptr = do
    width <- #{peek XImage, width} ptr
    height <- #{peek XImage, height} ptr
    xOffset <- #{peek XImage, xoffset} ptr
    format <- #{peek XImage, format} ptr
    imageData <- #{peek XImage, data} ptr
    byteOrder <- #{peek XImage, byte_order} ptr
    imageBitmapUnit <- #{peek XImage, bitmap_unit} ptr
    imageBitmapBitOrder <- #{peek XImage, bitmap_bit_order} ptr
    imageBitmapPad <- #{peek XImage, bitmap_pad} ptr
    depth <- #{peek XImage, depth} ptr
    bytesPerLine <- #{peek XImage, bytes_per_line} ptr
    bitsPerPixel <- #{peek XImage, bits_per_pixel} ptr
    redMask <- #{peek XImage, red_mask} ptr
    greenMask <- #{peek XImage, green_mask} ptr
    blueMask <- #{peek XImage, blue_mask} ptr
    return Image { width, height, xOffset, format, imageData, byteOrder, imageBitmapUnit, imageBitmapBitOrder, imageBitmapPad, depth, bytesPerLine, bitsPerPixel, redMask, greenMask, blueMask }
  poke ptr Image { width, height, xOffset, format, imageData, byteOrder, imageBitmapUnit, imageBitmapBitOrder, imageBitmapPad, depth, bytesPerLine, bitsPerPixel, redMask, greenMask, blueMask } = do
    #{poke XImage, width} ptr width
    #{poke XImage, height} ptr height
    #{poke XImage, xoffset} ptr xOffset
    #{poke XImage, format} ptr format
    #{poke XImage, data} ptr imageData
    #{poke XImage, byte_order} ptr byteOrder
    #{poke XImage, bitmap_unit} ptr imageBitmapUnit
    #{poke XImage, bitmap_bit_order} ptr imageBitmapBitOrder
    #{poke XImage, bitmap_pad} ptr imageBitmapPad
    #{poke XImage, depth} ptr depth
    #{poke XImage, bytes_per_line} ptr bytesPerLine
    #{poke XImage, bits_per_pixel} ptr bitsPerPixel
    #{poke XImage, red_mask} ptr redMask
    #{poke XImage, green_mask} ptr greenMask
    #{poke XImage, blue_mask} ptr blueMask

getImage :: Display -> Drawable -> Position -> Position -> Dimension -> Dimension -> PlaneMask -> ImageFormat -> IO (Ptr Image)
getImage display d x y width height plane_mask format = do
  imagePtr <- xGetImage display d x y width height plane_mask format
  when (imagePtr == nullPtr) $ failMemory "getImage"
  return imagePtr

foreign import ccall unsafe "XGetImage" xGetImage :: Display -> Drawable -> Position -> Position -> Dimension -> Dimension -> PlaneMask -> ImageFormat -> IO (Ptr Image)

foreign import ccall unsafe "XGetPixel" getPixel :: Ptr Image -> Position -> Position -> IO Pixel

foreign import ccall unsafe "XWriteBitmapFile" writeBitmapFile :: Display -> CString -> Pixmap -> Dimension -> Dimension -> Position -> Position -> IO ()

foreign import ccall unsafe "XGetHeader" getHeader :: Ptr WindowAttributes -> Ptr Image -> CInt -> Ptr a -> CInt -> IO ()

foreign import ccall unsafe "XReadColors" readColors :: Display -> Ptr WindowAttributes -> Ptr (Ptr Color) -> IO CInt

foreign import ccall unsafe "XGetXWDColor" getXwdColor :: Ptr (Ptr Color) -> Ptr a -> CInt -> IO ()

foreign import ccall unsafe "XFreeReadColors" freeReadColors :: Ptr (Ptr Color) -> IO ()

saveImage :: Display -> Window -> Ptr Image -> Handle -> IO ()
saveImage display window imagePtr h = do
  image <- peek imagePtr
  let bufferSize = imageSize image
      headerSize = #{const SIZEOF(XWDheader)} + #{size char} -- size of char is included for the null string terminator of otherwise empty window filename
  alloca $ \windowAttributesPtr -> do
    _ <- xGetWindowAttributes display window windowAttributesPtr -- TODO: Check return value
    alloca $ \colorsPtr ->
      bracket (readColors display windowAttributesPtr colorsPtr) (\ncolors -> when (ncolors > 0) $ freeReadColors colorsPtr) $ \nColors -> do
        when (nColors < 0) $ failMemory "readColors"
        allocaBytes headerSize $ \headerPtr -> do
          -- TODO: Should we create header once and use it for every image, provided size and other attributes do not change? Do we benefit from that?
          getHeader windowAttributesPtr imagePtr nColors headerPtr (fromIntegral headerSize)
          hPutBuf h headerPtr headerSize
        allocaBytes #{const SIZEOF(XWDColor)} $ \xwdColorPtr ->
          forM_ [0..nColors-1] $ \i -> do
            getXwdColor colorsPtr xwdColorPtr i
            hPutBuf h xwdColorPtr #{const SIZEOF(XWDColor)}
  hPutBuf h (imageData image) bufferSize
    where imageSize image | format image /= zPixmap = fromIntegral (bytesPerLine image) * fromIntegral (height image) * fromIntegral (depth image)
                          | otherwise               = fromIntegral (bytesPerLine image) * fromIntegral (height image)

failMemory :: String -> IO ()
failMemory name = throwIO $ IOError Nothing ResourceExhausted name "out of memory" Nothing Nothing