{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

module Python(
    module Str,
    defO,
    defV,
    defOO,
    defOV,
    defVO,
    defVV,
    defOOO,
    defOOV,
    defOVO,
    defOVV,
    defVOO,
    defVOV,
    defVVO,
    defVVV,
    defOOOO,
    PyObject,
    )
where

import Str
import Data.Digest.Pure.MD5 (md5)
import Foreign.C
import Foreign (FunPtr, ForeignPtr)
import Foreign.Ptr (Ptr, nullPtr)
import Foreign.ForeignPtr (newForeignPtr, withForeignPtr)
import Data.Aeson (FromJSON, ToJSON, toJSON, encode, decode)
import Data.ByteString.Lazy.Char8 (unpack, pack)
import Data.IORef (IORef, newIORef, writeIORef, readIORef)
import System.IO.Unsafe (unsafePerformIO)
import qualified Data.Map as Map
import Data.Maybe (fromJust)

modules :: IORef (Map.Map String RawPyObject)
modules = unsafePerformIO $ newIORef (Map.empty)

data P
type RawPyObject = (Ptr P)
data PyObject a = PyObject (ForeignPtr P) deriving Show
    

foreign import ccall "getObject" c_getObject:: CString -> IO (RawPyObject)
foreign import ccall "getObjectInModule" c_getObjectInModule :: CString -> CString -> IO (RawPyObject)
foreign import ccall "execInModule" c_execInModule :: CString -> CString -> IO ()
foreign import ccall "PyRun_SimpleString" pyRun_SimpleString :: CString -> IO ()
foreign import ccall "Py_DecRef" py_DecRef :: RawPyObject -> IO ()
foreign import ccall "Py_Initialize" py_initialize :: IO ()
foreign import ccall "Py_Finalize" finalize :: IO ()
foreign import ccall "Py_BuildValue" py_BuildValueString :: CString -> CString -> IO (RawPyObject)
foreign import ccall "Py_BuildValue" py_BuildValueString2 :: CString -> CString -> CString -> IO (RawPyObject)
foreign import ccall "Py_BuildValue" py_BuildValueObject :: CString -> RawPyObject -> IO (RawPyObject)
foreign import ccall "Py_BuildValue" py_BuildValueObject2 :: CString -> RawPyObject -> RawPyObject -> IO (RawPyObject)
foreign import ccall "Py_BuildValue" py_BuildValueObject3 :: CString -> RawPyObject -> RawPyObject -> RawPyObject -> IO (RawPyObject)
foreign import ccall "PyObject_CallObject" pyObject_CallObject :: RawPyObject -> RawPyObject -> IO (RawPyObject)
foreign import ccall "PyString_AsString" pyString_AsString :: RawPyObject -> IO CString
foreign import ccall unsafe "gimmeFunc" gimmeFunc :: CInt -> IO (FunPtr (RawPyObject -> IO ()))

exec :: String -> IO ()
exec s = withCString s pyRun_SimpleString

withCString2 :: String -> String -> (CString -> CString -> IO a) -> IO a
withCString2 s1 s2 f = withCString s1 f' where
    f' cs1 = withCString s2 (f cs1)

execInModule :: String -> String -> IO ()
execInModule moduleName payload = 
    withCString2 payload moduleName c_execInModule

initialize :: IO ()
initialize = do
  py_initialize

getObjectInModule :: String -> String -> IO (RawPyObject)
getObjectInModule moduleName objectName = 
    withCString2 objectName moduleName c_getObjectInModule

getObject :: String -> IO (RawPyObject)
getObject s = withCString s c_getObject

jsonfunc :: String
jsonfunc = [str|
def jsonfunc(argformats):
    def wrapper(f):
        import simplejson as json
        import traceback
        def new_f(*args):
            new_args = []
            for (x, format) in zip(args, argformats):
                if format == 'J':
                    try:
                        new_x = json.loads(x)
                    except:
                        print('error processing %s' % x)  
                    new_args.append(new_x)
                else:
                    new_args.append(x)
            try:
                result = f(*new_args)
                if argformats[-1] == 'J':
                    result = json.dumps(result)
            except Exception as ex:
                print(traceback.format_exc())
                result = None
                if argformats[-1] == 'J':
                    result = json.dumps(result)
            return result
        return new_f
    return wrapper
|]

hash :: String -> String
hash contents = show . md5 $ pack contents

mydecode :: (FromJSON a) => String -> Maybe a
mydecode s = do
    x <- (decode . pack . (\x -> "[" ++ x ++ "]") ) s
    return $ head x


toPyObject :: (ToJSON a) => a -> IO (PyObject b)
toPyObject x = do
    y <- return (unpack . encode $ x)
    p <- withCString "s" (\cs -> 
           withCString y (\cy -> 
             py_BuildValueString cs cy) )
    newForeignPyPtr p

fromPyObject :: (FromJSON a) => PyObject b -> IO a
fromPyObject (PyObject fr) = do
    r2 <- withForeignPtr fr $ \r -> peekCString =<< pyString_AsString r
    return $ fromJust $ mydecode r2

getFunc :: String -> IO RawPyObject
getFunc s = do
    currentModules <- readIORef $ modules
    key <- return $ hash s
    if Map.member "initialized" currentModules then
      return ()
    else initialize
    case Map.lookup key currentModules of
      Just p -> return p
      Nothing -> do 
        execInModule key jsonfunc
        execInModule key s
        execInModule key "export = jsonfunc('JJ')(export)"
        f' <- getObjectInModule key "export"
        writeIORef modules (Map.insert (hash s) f' currentModules)
        return f'

newForeignPyPtr :: RawPyObject -> IO (PyObject b)
newForeignPyPtr r = do
    finalizer <- gimmeFunc 0
    return . PyObject =<< newForeignPtr finalizer r

defO :: String -> IO (PyObject b)
defO s = do
    f <- getFunc s
    r <- pyObject_CallObject f nullPtr
    newForeignPyPtr r

defOO :: String -> (PyObject a) -> IO (PyObject b)
defOO s (PyObject fx1) = do
    f <- getFunc s
    p1 <- withForeignPtr fx1 $ \x1 -> withCString "(O)" (\cs -> py_BuildValueObject cs x1) 
    (PyObject fp) <- newForeignPyPtr p1
    r <- withForeignPtr fp $ \p -> pyObject_CallObject f p
    newForeignPyPtr r

defOOO :: String -> (PyObject a1) -> (PyObject a2) -> IO (PyObject b)
defOOO s (PyObject fx1) (PyObject fx2) = do
    f <- getFunc s
    p1 <- withForeignPtr fx1 $ \x1 -> (
        withForeignPtr fx2 $ \x2 -> (
            withCString "(OO)" (\cs -> py_BuildValueObject2 cs x1 x2) 
        ))
    (PyObject fp) <- newForeignPyPtr p1
    r <- withForeignPtr fp $ \p -> pyObject_CallObject f p
    newForeignPyPtr r

defOOOO :: String -> (PyObject a1) -> (PyObject a2) -> (PyObject a3) -> IO (PyObject b)
defOOOO s (PyObject fx1) (PyObject fx2) (PyObject fx3) = do
    f <- getFunc s
    p1 <- withForeignPtr fx1 $ \x1 -> (
        withForeignPtr fx2 $ \x2 -> (
            withForeignPtr fx3 $ \x3 -> (
                withCString "(OOO)" (\cs -> py_BuildValueObject3 cs x1 x2 x3) 
        )))
    (PyObject fp) <- newForeignPyPtr p1
    r <- withForeignPtr fp $ \p -> pyObject_CallObject f p
    newForeignPyPtr r

defV :: (FromJSON b) => String -> IO b
defV s = do
    r <- defO s
    b <- fromPyObject r
    return b

defOV :: (FromJSON b) => String -> (PyObject a1) -> IO b
defOV s x1 = do
    r <- defOO s x1
    b <- fromPyObject r
    return b

defVO :: (ToJSON a1) => String -> a1 -> IO (PyObject b)
defVO s input1 = do
    x1 <- toPyObject input1
    r <- defOO s x1
    return r

defVV :: (ToJSON a1, FromJSON b) => String -> a1 -> IO b
defVV s input1 = do
    x1 <- toPyObject input1
    fr <- defOO s x1
    b <- fromPyObject fr
    return b

defOOV :: (FromJSON b) => String -> (PyObject a1) -> (PyObject a2) -> IO b
defOOV s x1 x2 = do
    fr <- defOOO s x1 x2
    b <- fromPyObject fr
    return b

defOVO :: (ToJSON a2) => String -> (PyObject a1) -> a2 -> IO (PyObject b)
defOVO s x1 input2 = do
    x2 <- toPyObject input2
    fr <- defOOO s x1 x2
    return fr

defOVV :: (ToJSON a2, FromJSON b) => String -> (PyObject a1) -> a2 -> IO b
defOVV s x1 input2 = do
    x2 <- toPyObject input2
    fr <- defOOO s x1 x2
    b <- fromPyObject fr
    return b

defVOO :: (ToJSON a1) => String -> a1 -> (PyObject a2) -> IO (PyObject b)
defVOO s input1 x2 = do
    x1 <- toPyObject input1
    fr <- defOOO s x1 x2
    return fr

defVOV :: (ToJSON a1, FromJSON b) => String -> a1 -> (PyObject a2) -> IO b
defVOV s input1 x2 = do
    x1 <- toPyObject input1
    fr <- defOOO s x1 x2
    b <- fromPyObject fr
    return b

defVVO :: (ToJSON a1, ToJSON a2) => String -> a1 -> a2 -> IO (PyObject b)
defVVO s input1 input2 = do
    x1 <- toPyObject input1
    x2 <- toPyObject input2
    fr <- defOOO s x1 x2
    return fr

defVVV :: (ToJSON a1, ToJSON a2, FromJSON b) => String -> a1 -> a2 -> IO b
defVVV s input1 input2 = do
    x1 <- toPyObject input1
    x2 <- toPyObject input2
    fr <- defOOO s x1 x2
    b <- fromPyObject fr
    return b