{-# LANGUAGE DeriveDataTypeable, MultiParamTypeClasses, UndecidableInstances, FlexibleInstances, FlexibleContexts, ScopedTypeVariables, TypeFamilies, ImplicitParams, TypeOperators, ForeignFunctionInterface #-}
{-# LANGUAGE Trustworthy #-}

-- | Serialization of object graphs (pointer swizzling). 'readGraph' and 'writeGraph' are basic operations. You can also locate specific objects to read, then read them using 'readGraphAt'. Objects can be located using the low-level pointer operations from File.Mapped.
--
--   Possible errors:
--
--     * /Wrong file type/ - A program tried to read a file that wasn't created with 'openMapped'/'checkedOpen'/
--       'writeGraph'.
--
--     * /Structural mismatch/ - The file was created with one of those functions, but the type information
--       saved with the file doesn't match the type used with 'readGraph'. This error would also happen
--       if the file was created with low-level functions only.
--
--     * /mmap of ... failed/ - Some data from the file could not be brought into memory, or new space could not
--       be made in the file. Use 'getErrno' to get further information.
module File.Graph (SwizzleState, SwizM, UnswizzleState, UnswizM, Swiz(..), defUnswizzle, defSwizzle, writeGraphAt, readGraphAt, checkedOpen, writeGraph, readGraph, clone, ffirst, fsecond, elementAt) where

import Data.Dynamic
import Data.IORef
import Data.Int
import Data.Char
import Data.Word
import Data.List
import Data.Maybe
import Data.Endian
import qualified Data.Map as M
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as BC
import Control.Concurrent.MVar
import Control.Exception
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.StablePtr
import Foreign.Storable
import qualified Foreign.Marshal.Alloc as A
import System.Directory
import System.Mem
import Control.Monad.State.Lazy
import Generics.Pointless.Functors
import Data.WeakDict
import File.Mapped
import Data.Storables

import System.FileLock

type instance Represent (IORef t) = FilePtr t
type instance Represent (MVar t) = FilePtr t
type instance Represent (ForeignPtr t) = FilePtr t
type instance Represent (StablePtr t) = t
type instance Represent (Maybe t) = Maybe (Represent t)
type instance Represent [t] = FilePtr (Int32, t) -- pointers to the size and the first element
type instance Represent (Either t u) = Either (Represent t) (Represent u)
type instance Represent B.ByteString = FilePtr (Int32, Word8) -- pointers to the size and the first byte
type instance Represent (Fix f) = FilePtr (Rep f (Fix f))
type SwizzleState = M.Map Pointer Dyn

type SwizM t = StateT SwizzleState IO t

type UnswizzleState = WeakDict Dyn Pointer

type UnswizM t = StateT (UnswizzleState, UnswizzleState) IO t

-- | The class of objects that can be swizzled/unswizzled.
class Swiz t where
	size :: t -> Int
	aligned :: t -> Int
	unswizzle :: FilePtr t -> t -> UnswizM ()
	swizzle :: FilePtr t -> SwizM t

err = error "File.Graph.checkedOpen: structural mismatch"

defUnswizzle :: forall r t. (Eq (r t), Represent (r t) ~ FilePtr t, Swiz t, Typeable (r t)) => (r t -> IO t) -> FilePtr (r t) -> r t -> UnswizM ()
defUnswizzle read p ref = do
	(tabu, parm) <- get
	lift (lookupWeak (makeDyn ref) tabu) >>= maybe
		(lift (lookupWeak (makeDyn ref) parm) >>= maybe
			(do
				tabu' <- lift $ insertWeak (makeDyn ref) (offset p) tabu
				put (tabu', parm)
				lift $ malloc (fileSrc p) (size (undefined :: t)))
			(\o -> do
				tabu' <- lift $ insertWeak (makeDyn ref) o tabu
				put (tabu', parm)
				lift $ make (fileSrc p) o) >>= \rp -> do
			lift (read ref) >>= unswizzle rp
			lift $ pokePtr (toRepr p) rp)
		(lift . fpoke (toRepr p))

defSwizzle :: (Eq (r t), Represent (r t) ~ FilePtr t, Swiz t, Typeable (r t)) => IO (r t) -> (r t -> t -> IO ()) -> FilePtr (r t) -> SwizM (r t)
defSwizzle create write p = do
	st <- get
	o <- lift $ fpeek $ toRepr p
	maybe
		(do
			ref <- lift create
			put $ M.insert o (makeDyn ref) st
			rp <- lift $ peekPtr $ toRepr p
			swizzle rp >>= lift . write ref
			return ref)
		(return . fromJust . value)
		(M.lookup o st)

-- The implementation uses equality testing on references, to break cycles.
instance (Typeable t, Swiz t) => Swiz (IORef t) where
	size _ = 4
	aligned _ = 4
	unswizzle = defUnswizzle readIORef
	swizzle = defSwizzle (newIORef undefined) writeIORef

instance (Typeable t, Swiz t) => Swiz (MVar t) where
	size _ = 4
	aligned _ = 4
	unswizzle = defUnswizzle readMVar
	swizzle = defSwizzle newEmptyMVar putMVar

type CFunction t = Ptr t -> IO ()

foreign import ccall "wrapper" makeFin :: CFunction t -> IO (FunPtr (CFunction t))

instance (Typeable t, Storable t, Swiz t) => Swiz (ForeignPtr t) where
	size _ = 4
	aligned _ = 4
	unswizzle = defUnswizzle (`withForeignPtr` peek)
	swizzle = defSwizzle (makeFin A.free >>= \fr -> A.malloc >>= newForeignPtr fr) (\ptr x -> withForeignPtr ptr (`poke` x))

-- StablePtrs are not mutable, therefore they save an indirection vs. IORef
instance (Swiz t, Typeable t) => Swiz (StablePtr t) where
	size _ = size (undefined :: t)
	aligned _ = aligned (undefined :: t)
	unswizzle p sp = do
		(tabu, parm) <- get
		lift (lookupWeak (makeDyn sp) tabu) >>= maybe
			(do
				x <- lift $ deRefStablePtr sp
				tabu' <- lift $ insertWeak (makeDyn sp) (offset p) tabu
				put (tabu', parm)
				unswizzle (toRepr p) x)
			(\_ -> return ())
	swizzle p = do
		st <- get
		maybe
			(do
				x <- swizzle $ toRepr p
				sp <- lift $ newStablePtr x
				put $ M.insert (offset p) (makeDyn sp) st
				return sp)
			(return . fromJust . value)
			(M.lookup (offset p) st)

instance (Swiz t) => Swiz (Maybe t) where
	size _ = size (undefined :: t) + 1
	aligned _ = aligned (undefined :: t)
	unswizzle p may = unswizzle (coerce p) $ maybe (Left ()) Right may
	swizzle p = liftM (either (\_ -> Nothing) Just) $ swizzle (coerce p :: FilePtr (Either () t))

step :: forall t. (Swiz t) => t -> Int
step _ = size (undefined :: t) `lcm` aligned (undefined :: t)

instance (Swiz t) => Swiz [t] where
	size _ = 4
	aligned _ = 4
	unswizzle p ls = do
		p2 <- lift $ malloc (fileSrc p) (4 `lcm` aligned (undefined :: t) + step (undefined :: t) * length ls)
		lift $ fpoke (ffirst p2) $ fromIntegral $ length ls
		let p3 = fsecond p2
		mapM_ (\(i, x) -> unswizzle (plus p3 i) x) $ zip [0,step (undefined :: t)..] ls
		lift $ pokePtr (toRepr p) p2
	swizzle p = do
		p2 <- lift $ peekPtr $ toRepr p
		m <- lift $ fpeek $ ffirst p2
		let p3 = fsecond p2
		mapM (swizzle . plus p3) [0,step (undefined :: t)..step (undefined :: t)*(fromIntegral m-1)]

instance (Swiz t, Swiz u) => Swiz (Either t u) where
	size ei = size (undefined :: t) `max` size (undefined :: u) + 1
	aligned _ = aligned (undefined :: t)
		`lcm` aligned (undefined :: u)
	unswizzle p ei = do
		lift $ poke (plusPtr (castPtr (pointer p)) (size (undefined :: t) `max` size (undefined :: u))) (either (\_ -> 0) (\_ -> 1) ei :: Int8)
		either
			(unswizzle (coerce p))
			(unswizzle (coerce p))
			ei
	swizzle p = do
		by :: Int8 <- lift $ peek $ plusPtr (castPtr (pointer p)) (size (undefined :: t) `max` size (undefined :: u))
		if by == 0 then
				liftM Left $ swizzle $ coerce p
			else
				liftM Right $ swizzle $ coerce p

instance (Swiz t, Swiz u) => Swiz (t, u) where
	size _ = align' (undefined :: t) (undefined :: u) + size (undefined :: u)
	aligned  _ = aligned (undefined :: t)
		`lcm` aligned  (undefined :: u)
	unswizzle p (x, y) = do
		unswizzle (ffirst p) x
		unswizzle (fsecond p) y
	swizzle p = liftM2 (,) (swizzle (ffirst p)) (swizzle (fsecond p))

instance Swiz () where
	size = sizeOf
	aligned = alignment
	unswizzle _ () = return ()
	swizzle _ = return ()

instance Swiz Int8 where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Int16 where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Int32 where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p . fromIntegral
	swizzle = lift . fpeek

instance Swiz Int where
	size = sizeOf
	aligned = alignment
	unswizzle p x = lift $ fpoke p (fromIntegral x :: Int32)
	swizzle p = lift $ liftM (\(x :: Int32) -> fromIntegral x) $ fpeek p

instance Swiz Int64 where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Word8 where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Word16 where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Pointer where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Word64 where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Float where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Char where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz Bool where
	size = sizeOf
	aligned = alignment
	unswizzle p = lift . fpoke p
	swizzle = lift . fpeek

instance Swiz B.ByteString where
	size _ = 4
	aligned _ = 4
	unswizzle p bs = lift $ do
		p2 <- malloc (fileSrc p) (4 + B.length bs)
		poke (first (pointer p2)) $ fromIntegral $ B.length bs
		let p3 = second $ pointer p2
		mapM_ (\i -> poke (plusPtr p3 i) (B.index bs i)) [0..B.length bs-1]
		pokePtr (toRepr p) p2
	swizzle p = lift $ do
		p2 <- peekPtr (toRepr p)
		m <- peek $ first $ pointer p2
		let p3 = second $ pointer p2
		ls <- mapM (peek . plusPtr p3) [0..fromIntegral m-1]
		return (B.pack ls)

-- The fixpoint constructor allows one to serialize inductive data types.
instance (Swiz (Rep f (Fix f))) => Swiz (Fix f) where
	size _ = 4
	aligned _ = 4
	unswizzle p (Inn f) = do
		p2 <- lift $ malloc (fileSrc p) (size f)
		unswizzle p2 f
		lift $ pokePtr (toRepr p) p2
	swizzle p = do
		p2 <- lift $ peekPtr (toRepr p)
		liftM Inn $ swizzle p2

instance Swiz (FilePtr t) where
	size _ = 4
	aligned _ = 4
	-- This is only meaningful if the two pointers point into the same file.
	unswizzle p = lift . pokePtr p
	swizzle = lift . peekPtr

instance Storable (FilePtr t) where
	sizeOf _ = 4
	alignment _ = 4
	poke p = poke (castPtr p) . offset
	peek = error "File.Graph.peek: not supported for FilePtr"

instance EndianSensitive (FilePtr t) where
	swapEndian = id

convertDict = foldM (\dict (o, dyn) -> insertWeak dyn o dict)

-- | Save an object graph into the location of the FilePtr.
--
--   A "SwizzleState" parameter produced by a previous call to 'readGraphAt' can be supplied.
--   This will cause the references created by that call to be written back to the same
--   locations in the file. If this is not needed, an empty SwizzleState can be supplied,
--   in which case fresh references are created in the file.
writeGraphAt :: (Swiz t) =>  FilePtr t -> t -> SwizzleState -> IO ()
writeGraphAt p x parm = do
	tabu <- emptyWeak
	newDict <- emptyWeak
	parm' <- convertDict newDict $ M.assocs parm
	void $ runStateT (unswizzle p x) (tabu, parm')

-- | Read an object graph from the location.
readGraphAt p = runStateT (swizzle p) M.empty

-- | Obtain a root pointer, but checking the type-integrity of the pointer.
--
--   The root pointer points to the second DWORD of the root object.
checkedOpen :: forall t. (Swiz t, Typeable t) => FilePath -> IO (FilePtr t)
checkedOpen path = do
	p <- openMapped path
	(ty, _) <- readGraphAt $ ffirst p
	unless (ty == BC.pack (show $ typeOf (undefined :: t))) err
	peekPtr $ fsecond p

-- | Save an object graph into the file.
writeGraph path x state = do
	p <- openMapped path
	finally
		(do
		let ty = BC.pack $ show $ typeOf x
		writeGraphAt (ffirst p) ty M.empty
		p2 <- malloc (fileSrc p) (size x)
		pokePtr (fsecond p) p2
		writeGraphAt p2 x state)
		(closeMapped (fileSrc p))

-- | Read an object graph from the file.
readGraph path = do
	p <- checkedOpen path
	finally (readGraphAt p) (closeMapped (fileSrc p))

-- | Create a copy of the parameter with fresh references.
clone :: (Swiz t) => t -> IO t
clone x = do
	dir <- getTemporaryDirectory
	let tmp = dir ++ "GRAPH_CLONE.tmp"
	maskSigint $ withFileLock (dir ++ "GRAPH_CLONE_1.tmp.lock") Exclusive $ const $ do
		p <- openMapped tmp
		finally
			(do
			p <- malloc (fileSrc p) (size x)
			writeGraphAt p x M.empty
			(y, _) <- readGraphAt p
			return y)
			(do
			closeMapped (fileSrc p)
			removeFile tmp)

test :: Int -> Fix (Const () :+: ((Const Int :*: Const String) :*: Id))
test 0 = Inn (Left ())
test n = Inn (Right ((n, "Test of stuff"), test (n - 1)))

ffirst :: FilePtr (t, u) -> FilePtr t
ffirst = coerce

align' x y = ((size x - 1) `div` aligned y + 1) * aligned y

fsecond :: forall a b. (Swiz a, Swiz b) => FilePtr (a, b) -> FilePtr b
fsecond = coerce . (`plus` fromIntegral (align' (undefined :: a) (undefined :: b)))

elementAt :: forall t. (Swiz t) => FilePtr [t] -> Int32 -> IO (FilePtr t)
elementAt fp i = do
	p <- peekPtr (toRepr fp)
	len <- fpeek (ffirst p)
	return $ if i < 0 || i >= len then
			error "File.Graph.elementAt: index out of range"
		else
			plus (fsecond p) (fromIntegral i * size (undefined :: t))