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

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

data EnvVar = NotSet | Literal String | forall a.Show a => Scalar a | Array [EnvVar]

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

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')

localEnvironment :: Functor m => ExpanderT m a -> m a
localEnvironment m = fmap fst $ runExpanderT m []

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

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

class Monad ee => ExpanderEnv ee where
  mgetv :: String -> ee EnvVar
  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

class Monad e => MonadExpand e where
  expand :: String -> e String

instance MonadExpand IO where
  expand = expandVars

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

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')

isAnum = (`elem` (['a'..'z']++['A'..'Z']++"_"++['0'..'9']))