{-# LANGUAGE DefaultSignatures    #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE InstanceSigs         #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

module Foreign.Storable.ATS
    ( ATSStorable (..)
    ) where

import           Control.Composition
import           Data.Word
import           Foreign.Marshal.Alloc
import           Foreign.Ptr
import qualified Foreign.Storable      as C
import           GHC.Generics

class Storable' f where

    sizeOf' :: f a -> Int

    alignment' :: f a -> Int

    peek' :: Ptr (f a) -> IO (f a)

    poke' :: Ptr (f a) -> f a -> IO ()

    pokeByteOff' :: Ptr (f a) -> Int -> f a -> IO ()
    pokeByteOff' = poke' .* plusPtr

    peekByteOff' :: Ptr (f a) -> Int -> IO (f a)
    peekByteOff' = peek' .* plusPtr

instance Storable' V1 where
    peek' = undefined
    poke' = undefined
    alignment' = undefined
    sizeOf' = undefined

-- TODO closed type families?

instance (Storable' a, Storable' b) => Storable' (a :*: b) where
    sizeOf' _ = sizeOf' (undefined :: a x) + sizeOf' (undefined :: b x)
    alignment' _ = gcd (alignment' (undefined :: a x)) (alignment' (undefined :: b x))
    peek' ptr = do
        a <- peek' (castPtr ptr)
        (a :*:) <$> peekByteOff' (castPtr ptr) (sizeOf' a)
    poke' ptr (a :*: b) = mconcat
        [ poke' (castPtr ptr) a
        , pokeByteOff' (castPtr ptr) (sizeOf' a) b ]

instance C.Storable a => ATSStorable a where
    sizeOf :: a -> Int
    sizeOf = C.sizeOf
    alignment :: a -> Int
    alignment = C.alignment
    poke :: Ptr a -> a -> IO ()
    poke = C.poke
    peek :: Ptr a -> IO a
    peek = C.peek

instance (Storable' a, Storable' b) => Storable' (a :+: b) where
    sizeOf' _ = 4 + max (sizeOf' (undefined :: a x)) (sizeOf' (undefined :: b x))
    alignment' _ = gcd 4 $ gcd (alignment' (undefined :: a x)) (alignment' (undefined :: b x))
    peek' ptr = do
        tag <- peek (castPtr ptr)
        if (tag :: Word32) == 0
            then pure L1 <*> peekByteOff' (castPtr ptr) 4
            else pure R1 <*> peekByteOff' (castPtr ptr) 4
    poke' ptr (L1 val) = mconcat
        [ poke (castPtr ptr) (0 :: Word32)
        , pokeByteOff' (castPtr ptr) 4 val ]
    poke' ptr (R1 val) = mconcat
        [ poke (castPtr ptr) (1 :: Word32)
        , pokeByteOff' (castPtr ptr) 4 val ]

instance (ATSStorable a) => Storable' (K1 i a) where
    sizeOf' _ = sizeOf (undefined :: a)
    alignment' _ = alignment (undefined :: a)
    peek' ptr = pure K1 <*> peek (castPtr ptr)
    poke' ptr (K1 val) = poke (castPtr ptr) val

instance (Storable' a) => Storable' (M1 i c a) where
    sizeOf' _ = sizeOf' (undefined :: a x)
    alignment' _ = alignment' (undefined :: a x)
    peek' ptr = pure M1 <*> peek' (castPtr ptr)
    poke' ptr (M1 val) = poke' (castPtr ptr) val

class ATSStorable a where

    sizeOf :: a -> Int
    default sizeOf :: (Generic a, Storable' (Rep a)) => a -> Int
    sizeOf _ = (sizeOf' . from) (undefined :: a)

    alignment :: a -> Int
    default alignment :: (Generic a, Storable' (Rep a)) => a -> Int
    alignment _ = (alignment' . from) (undefined :: a)

    poke :: Ptr a -> a -> IO ()
    default poke :: (Generic a, Storable' (Rep a)) => Ptr a -> a -> IO ()
    poke ptr x = poke' (castPtr ptr) (from x)

    peek :: Ptr a -> IO a
    default peek :: (Generic a, Storable' (Rep a)) => Ptr a -> IO a
    peek = fmap to . peek' . castPtr

    writePtr :: a -> IO (Ptr a)
    writePtr val = do
        ptr <- mallocBytes (sizeOf val)
        poke ptr val
        pure ptr