{-# LANGUAGE ExistentialQuantification, RankNTypes, Rank2Types #-}

-- | Provides generic string expansion and a variable expander
module Text.Chatty.Expansion where

import Data.List
import Control.Arrow
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import System.Environment hiding (getEnv)
import System.Posix.Env (getEnv, setEnv)
--import System.SetEnv

-- | Some environment variable
data EnvVar = NotSet                       -- ^ Not set.
            | Literal String               -- ^ An embeddable string.
            | forall a.Show a => Scalar a  -- ^ Something we can show.
            | Array [EnvVar]               -- ^ Array of that

instance Show EnvVar where
  show (Scalar s) = show s
  show (Literal s) = s
  show (Array ps) = unwords $ map show ps
  show NotSet = ""

-- | Environment storage and variable expander.
newtype ExpanderT m a = Expander {
    runExpanderT :: [(String,EnvVar)] -> m (a,[(String,EnvVar)])
  }

instance Monad m => Monad (ExpanderT m) where
  return a = Expander $ \vs -> return (a,vs)
  (Expander e) >>= f = Expander $ \vs -> do (a,vs') <- e vs; runExpanderT (f a) vs'

instance MonadTrans ExpanderT where
  lift m = Expander $ \vs -> do a <- m; return (a,vs)

instance MonadIO m => MonadIO (ExpanderT m) where
  liftIO = lift . liftIO

instance Monad m => Functor (ExpanderT m) where
  fmap f a = Expander $ \vs -> do (a',vs') <- runExpanderT a vs; return (f a',vs')

-- | Run this function inside a blank environment.
localEnvironment :: Functor m => ExpanderT m a -> m a
localEnvironment m = fmap fst $ runExpanderT m []

-- | Run this function in a locally modifiable, but not exported environment
forkEnvironment :: (Functor m,Monad m,MonadIO m) => ExpanderT m a -> m a
forkEnvironment m = do
  es <- liftIO getEnvironment
  fmap fst $ runExpanderT m $ fmap (second Literal) es

-- | Export this local environment.
exportAll :: (Monad m,MonadIO m) => ExpanderT m ()
exportAll = Expander $ \vs -> do
  liftIO $ forM_ vs $ \(k,v) -> setEnv k (show v) True
  return ((),vs)

-- | Typeclass for all environment storages.
class Monad ee => ExpanderEnv ee where
  -- | Get environment variable
  mgetv :: String -> ee EnvVar
  -- | Put environment variable
  mputv :: String -> EnvVar -> ee ()

instance Monad m => ExpanderEnv (ExpanderT m) where
  mgetv s = Expander $ \vs -> return $
    case filter ((==s).fst) vs of
      [] -> (NotSet,vs)
      ((_,v):_) -> (v,vs)
  mputv k v = Expander $ \vs -> return ((),(k,v):filter ((/=k).fst) vs)

instance ExpanderEnv IO where
  mgetv = fmap (\v -> case v of Nothing -> NotSet; Just v' -> Literal v') . getEnv
  mputv k v = setEnv k (show v) True

-- | Typeclass for all string-expanding monads.
class Monad e => MonadExpand e where
  -- | Expand the given string.
  expand :: String -> e String

instance MonadExpand IO where
  expand = expandVars

instance Monad m => MonadExpand (ExpanderT m) where
  expand = expandVars

-- | Expand $variables.
expandVars :: (Monad m,Functor m,ExpanderEnv m) => String -> m String
expandVars [] = return []
expandVars ('$':ss) =
  let nm = takeWhile isAnum ss
      rm = dropWhile isAnum ss
  in do
    v <- fmap show $ mgetv nm
    r <- expandVars rm
    return (v++r)
expandVars (s:ss) = do ss' <- expandVars ss; return (s:ss')

-- | Is alphanumeric?
isAnum = (`elem` (['a'..'z']++['A'..'Z']++"_"++['0'..'9']))