{-  CAO Compiler
    Copyright (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>. -}

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE BangPatterns #-}

{- |
Module      :  $Header$
Description :  General application monad.
Copyright   :  (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho
License     :  GPL

Maintainer  :  Paulo Silva <paufil@di.uminho.pt>
Stability   :  experimental
Portability :  non-portable

General application monad.

-}

module Language.CAO.Common.Monad 
    ( CaoMonad
    , CaoM(..)
    , CaoResult
    , CaoState
    , getFileName
    , setFileName
    , runCaoResultWarn
    , tcError
    , tcWarn
    , ensureDepMode
    , caoOrCalf
    , withStrictMode
    ) where

import Control.Applicative ( (<$>) )

import Control.Arrow ( first )

import Control.Monad.Error
import Control.Monad.State.Strict
import Control.Monad.Writer.Strict

import Language.CAO.Common.Error
import Language.CAO.Common.Outputable
import Language.CAO.Common.SrcLoc
import Language.CAO.Common.State
import Language.CAO.Common.Utils
import Language.CAO.Common.Var

import Main.Flags (RunMode(..))

class CaoM CaoError CaoWarning CaoState m => CaoMonad m where
instance CaoM CaoError CaoWarning CaoState m => CaoMonad m where

class (Functor m, Monad m, MonadIO m, MonadError e m, MonadWriter w m, MonadState s m) =>
    CaoM e w s m where
        uniqId          :: m Int
        injectResult    :: Either String a -> m a
        caoError        :: (Show id, Read id, PP id)
                              => SrcLoc -> ErrorCode id -> m a
        caoWarning      :: PP id => SrcLoc -> WarningCode id -> m ()

--------------------------------------------------------------------------------

newtype CaoResult a
    = CaoResult { runCaoResult :: 
        CaoState -> IO ( Either CaoError
                                (a, CaoWarning)
                       , CaoState) }

instance CaoM CaoError CaoWarning CaoState CaoResult where
    injectResult = either (throwError . read) return
    uniqId = getLastVar
    caoError info code = do
        fnm <- getFileName
        throwError $ mkCaoError info fnm code
    caoWarning info msg = do
        fnm <- getFileName
        tell $ mkCaoWarning $ mkCaoWarningInfo info fnm msg

instance Functor CaoResult where
    fmap f (CaoResult m) = CaoResult $! \ st -> first (either Left (Right . first f)) <$> m st

instance Monad CaoResult  where
    return !x = CaoResult $! \ st -> return (Right (x, mempty), st)
    (>>=)     = bindTcMonad

{-# INLINE bindTcMonad #-}
{-# INLINE bindTcMonad2 #-}
bindTcMonad :: CaoResult a
            -> (a -> CaoResult b)
            -> CaoResult b
bindTcMonad m f =  CaoResult $! \ st -> do
    (x', st') <- runCaoResult m st
    bindTcMonad2 x' (st', f)

bindTcMonad2 :: Either CaoError (a, CaoWarning)
             -> (CaoState, a -> CaoResult b)
             -> IO (Either CaoError (b, CaoWarning), CaoState)
bindTcMonad2 (Left !e) (!st',_) =
    return (Left e, st')
bindTcMonad2 (Right (!r, !w)) (!st', f) = do
    liftM (mapFst (fixR w)) $ runCaoResult (f r) st'
    where 
    fixR  _  !l@(Left _)        = l
    fixR !w' (Right (!x, !w'')) = Right (x, w' `mappend` w'')

instance MonadIO CaoResult where
    liftIO m = CaoResult $! \ st -> do
        r <- m
        return (Right (r, mempty), st) 

instance MonadState CaoState CaoResult where
    get     = CaoResult $! \ st -> return (Right (st, mempty), st)
    put !st = CaoResult $! \ _  -> return (Right ((), mempty), st)

instance MonadWriter CaoWarning CaoResult where
    tell   !w = CaoResult $! \ st -> return (Right ((), w), st)
    listen m  = CaoResult $! liftM (mapFst fixW) . runCaoResult m
        where 
        fixW :: Either CaoError (a, CaoWarning)
             -> Either CaoError ((a,CaoWarning), CaoWarning)
        fixW (Left !e)        = Left  e
        fixW (Right (!a, !w)) = Right ((a, w), w) 

    pass m  = CaoResult $! liftM (mapFst fixW) . runCaoResult m
        where 
        fixW :: Either CaoError ((a, CaoWarning -> CaoWarning), CaoWarning)
             -> Either CaoError (a, CaoWarning)
        fixW (Left !e)             = Left e
        fixW (Right ((!a, f), !w)) = Right (a, f w)

instance MonadError CaoError CaoResult where
    throwError !e   = CaoResult $! \ st -> return (Left e, st)
    catchError m f  = CaoResult $! go f <=< runCaoResult m
        where 
        go :: (CaoError -> CaoResult a)
           -> (Either CaoError (a, CaoWarning), CaoState)
           -> IO (Either CaoError (a, CaoWarning), CaoState)
        go c (Left e, st) = runCaoResult (c e) st
        go _ x = return x

runCaoResultT :: CaoResult a -> IO (Either CaoError (a, CaoWarning))
runCaoResultT m = liftM fst $ runCaoResult m initialState

runCaoResultWarn :: CaoResult a -> IO (a, CaoWarning)
runCaoResultWarn = either (fail . showCaoError) return <=< runCaoResultT

--------------------------------------------------------------------------------

tcError :: (CaoMonad m, PP id, Show id, Read id) => ErrorCode id -> m a
tcError err = do
    loc <- getSrcLoc
    caoError loc err

tcWarn :: (CaoMonad m, PP id) => WarningCode id -> m ()
tcWarn msg = do
    loc <- getSrcLoc
    caoWarning loc msg 

ensureDepMode :: CaoMonad m => m a -> m a
ensureDepMode t = do
    m <- getMode
    case m of
        CAO        -> tcError (StrictModeErr :: ErrorCode Var)
        CAO_Strict -> tcError (StrictModeErr :: ErrorCode Var)
        _          -> t

caoOrCalf :: CaoMonad m => m a -> m a -> m a
caoOrCalf cao calf = do
    m <- getMode
    case m of
        CAO        -> cao
        CAO_Strict -> cao
        _          -> calf

withStrictMode :: CaoMonad m => m a -> m a -> m a
withStrictMode tstrict tnstrict = do
    m <- getMode
    case m of
        CAO_Strict  -> tstrict
        CALF_Strict -> tstrict
        _           -> tnstrict