{-# LANGUAGE CPP #-}
module Data.STBImage.Immutable (Image(..), unsafeCastImage, loadImageBytes, writeNChannelPNG, writeNChannelBMP, writeNChannelTGA) where

import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV
import           Data.Either
import           Foreign
import           Foreign.C.Types
import           Foreign.C.String

import Data.STBImage.ColorTypes


-- | 'Image' is the least opinionated reasonable type to represent an image, just a vector of pixel 'Color's (laid out top-to-bottom, left-to-right) and a size.
data Image a = Image { _pixels :: V.Vector a, _width :: Int, _height :: Int }
           deriving (Eq) 

instance Show (Image a) where
    show (Image _ w h) = "Image (" ++ show w ++ "x" ++ show h ++ ")"

unsafeCastImage :: (Storable a, Storable b) => Image a -> Image b
unsafeCastImage (Image pixels w h) = Image (V.unsafeCast pixels) w h

--

foreign import ccall "stb_image.h stbi_load" stbi_load :: CString -> Ptr CInt -> Ptr CInt -> Ptr CInt -> CInt -> IO (Ptr CUChar)
foreign import ccall "stb_image.h stbi_failure_reason" stbi_failure_reason :: IO (CString)
foreign import ccall "stb_image.h &stbi_image_free" stbi_image_free :: FunPtr (Ptr CUChar -> IO ())

loadImageBytes :: Int -> FilePath -> IO (Either String (Image CUChar))
loadImageBytes comps path = do
    cPath <- newCString path
    widthPtr <- new 0
    heightPtr <- new 0
    nComponentsPtr <- new 0

    dataPtr <- stbi_load cPath widthPtr heightPtr nComponentsPtr (fromIntegral comps)

    case dataPtr /= nullPtr of
        True -> do
            dataForeignPtr <- newForeignPtr stbi_image_free dataPtr

            width  <- fromIntegral <$> peek widthPtr :: IO Int
            height <- fromIntegral <$> peek heightPtr :: IO Int

            let storage = V.unsafeFromForeignPtr0 dataForeignPtr (width * height * comps)

            free cPath
            free widthPtr
            free heightPtr
            free nComponentsPtr

            return $ Right (Image storage width height)
        False -> do
            err <- peekCString =<< stbi_failure_reason
            return $ Left err

--

foreign import ccall "stb/stb_image_write.h stbi_write_png" stbi_write_png :: CString -> CInt -> CInt -> CInt -> Ptr CUChar -> CInt -> IO CInt
foreign import ccall "stb/stb_image_write.h stbi_write_bmp" stbi_write_bmp :: CString -> CInt -> CInt -> CInt -> Ptr CUChar         -> IO CInt
foreign import ccall "stb/stb_image_write.h stbi_write_tga" stbi_write_tga :: CString -> CInt -> CInt -> CInt -> Ptr CUChar         -> IO CInt

writeNChannelPNG :: (Storable a) => CInt -> FilePath -> Image a -> IO ()
writeNChannelPNG comps path (Image storage width height) = do
    cPath <- newCString path

    let w = fromIntegral width :: CInt
    let h = fromIntegral height :: CInt

    withForeignPtr (fst $ V.unsafeToForeignPtr0 storage) (\pixBuf ->
        stbi_write_png cPath w h comps (castPtr pixBuf) (w * comps) -- bytes per row
        )

    free cPath

writeNChannelBMP :: (Storable a) => CInt -> FilePath -> Image a -> IO ()
writeNChannelBMP comps path (Image storage width height) = do
    cPath <- newCString path

    let w = fromIntegral width :: CInt
    let h = fromIntegral height :: CInt

    withForeignPtr (fst $ V.unsafeToForeignPtr0 storage) $ stbi_write_bmp cPath w h comps . castPtr

    free cPath


writeNChannelTGA :: (Storable a) => CInt -> FilePath -> Image a -> IO ()
writeNChannelTGA comps path (Image storage width height) = do
    cPath <- newCString path

    let w = fromIntegral width :: CInt
    let h = fromIntegral height :: CInt

    withForeignPtr (fst $ V.unsafeToForeignPtr0 storage) $ stbi_write_tga cPath w h comps . castPtr

    free cPath