{-# OPTIONS_HADDOCK not-home #-}

module Text.StringTemplate.Base
    (StringTemplate(..), StringTemplateShows(..), ToSElem(..), STGroup,
     Stringable(..), stShowsToSE, inSGen,
     toString, toPPDoc, render, newSTMP, newAngleSTMP,
     getStringTemplate, getStringTemplate',
     setAttribute, setManyAttrib,
     setNativeAttribute, setManyNativeAttrib,
     withContext, optInsertTmpl, setEncoder,
     paddedTrans, SEnv(..), parseSTMP, dumpAttribs,
     parseSTMPNames
    ) where
import Control.Monad
import Control.Arrow
import Control.Applicative hiding ((<|>),many)
import Data.Maybe
import Data.Monoid
import Data.List
import Text.ParserCombinators.Parsec
import qualified Data.Map as M
import qualified Text.PrettyPrint.HughesPJ as PP

import Text.StringTemplate.Classes
import Text.StringTemplate.Instances()

{--------------------------------------------------------------------
  Generic Utilities
--------------------------------------------------------------------}

type TmplParser x = GenParser Char ((Char, Char),[String]) x

(<$$>) :: (Functor f1, Functor f) => (a -> b) -> f (f1 a) -> f (f1 b)
(<$$>) x y = ((<$>) . (<$>)) x y
infixr 8 <$$>

(|.) :: (t1 -> t2) -> (t -> t1) -> t -> t2
(|.) f g x = f (g x)
infixr 3 |.

(.>>) :: (Monad m) => m a -> m b -> m b
(.>>) f g = f >> g
infixr 5 .>>

fromMany :: b -> ([a] -> b) -> [a] -> b
fromMany e _ [] = e
fromMany _ f xs  = f xs

swing :: (((a -> c1) -> c1) -> b -> c) -> b -> a -> c
swing = flip . (. flip id)

paddedTrans :: a -> [[a]] -> [[a]]
paddedTrans _ [] = []
paddedTrans n as = take (maximum . map length $ as) . trans $ as
    where trans ([] : xss)  = (n : map h xss) :  trans ([n] : (map t xss))
          trans ((x : xs) : xss) = (x : map h xss) : trans (m xs : (map t xss))
          trans _ = [];
          h (x:_) = x; h _ = n; t (_:y:xs) = (y:xs); t _ = [n];
          m (x:xs) = (x:xs); m _ = [n];

{--------------------------------------------------------------------
  StringTemplate and the API
--------------------------------------------------------------------}

-- | A function that generates StringTemplates.
-- This is conceptually a query function into a \"group\" of StringTemplates.
type STGroup a = String -> (StFirst (StringTemplate a))

-- | A String with \"holes\" in it. StringTemplates may be composed of any
-- 'Stringable' type, which at the moment includes 'String's, 'ByteString's,
-- PrettyPrinter 'Doc's, and 'Endo' 'String's, which are actually of type
-- 'ShowS'. When a StringTemplate is composed of a type, its internals are
-- as well, so it is, so to speak \"turtles all the way down.\"
data StringTemplate a = STMP {senv :: SEnv a,  runSTMP :: SEnv a -> a}

-- | Renders a StringTemplate to a String.
toString :: StringTemplate String -> String
toString = render

-- | Renders a StringTemplate to a 'Text.PrettyPrint.HughesPJ.Doc'.
toPPDoc :: StringTemplate PP.Doc -> PP.Doc
toPPDoc = render

-- | Generic render function for a StringTemplate of any type.
render :: Stringable a => StringTemplate a -> a
render = runSTMP <*> senv

-- | Parses a String to produce a StringTemplate, with \'$\'s as delimiters.
-- It is constructed with a stub group that cannot look up other templates.
newSTMP :: Stringable a => String -> StringTemplate a
newSTMP = STMP (SEnv M.empty [] mempty id) . parseSTMP ('$','$')

-- | Parses a String to produce a StringTemplate, delimited by angle brackets.
-- It is constructed with a stub group that cannot look up other templates.
newAngleSTMP :: Stringable a => String -> StringTemplate a
newAngleSTMP = STMP (SEnv M.empty [] mempty id) . parseSTMP ('<','>')

-- | Yields a StringTemplate with the appropriate attribute set.
-- If the attribute already exists, it is appended to a list.
setAttribute :: (ToSElem a, Stringable b) => String -> a -> StringTemplate b -> StringTemplate b
setAttribute s x st = st {senv = envInsApp s (toSElem x) (senv st)}

-- | Yields a StringTemplate with the appropriate attributes set.
-- If any attribute already exists, it is appended to a list.
setManyAttrib :: (ToSElem a, Stringable b) => [(String, a)] -> StringTemplate b -> StringTemplate b
setManyAttrib = flip . foldl' . flip $ uncurry setAttribute

-- | Yields a StringTemplate with the appropriate attribute set.
-- If the attribute already exists, it is appended to a list.
-- This will not translate the attribute through any intermediate
-- representation, so is more efficient when, e.g. setting
-- attributes that are large bytestrings in a bytestring template.
setNativeAttribute :: Stringable b => String -> b -> StringTemplate b -> StringTemplate b
setNativeAttribute s x st = st {senv = envInsApp s (SBLE x) (senv st)}

-- | Yields a StringTemplate with the appropriate attributes set.
-- If any attribute already exists, it is appended to a list.
-- Attributes are added natively, which may provide
-- efficiency gains.
setManyNativeAttrib :: (Stringable b) => [(String, b)] -> StringTemplate b -> StringTemplate b
setManyNativeAttrib = flip . foldl' . flip $ uncurry setNativeAttribute


-- | Replaces the attributes of a StringTemplate with those
-- described in the second argument. If the argument does not yield
-- a set of named attributes but only a single one, that attribute
-- is named, as a default, \"it\".
withContext :: (ToSElem a, Stringable b) => StringTemplate b -> a -> StringTemplate b
withContext st x = case toSElem x of
                     SM a -> st {senv = (senv st) {smp = a}}
                     b -> st {senv = (senv st) {smp = M.singleton "it" b}}

-- | Queries an String Template Group and returns Just the appropriate
-- StringTemplate if it exists, otherwise, Nothing.
getStringTemplate :: (Stringable a) => String -> STGroup a -> Maybe (StringTemplate a)
getStringTemplate s sg = stGetFirst (sg s)

-- | As with 'getStringTemplate' but never inlined, so appropriate for use
-- with volatile template groups.
{-# NOINLINE getStringTemplate' #-}
getStringTemplate' :: (Stringable a) => String -> STGroup a -> Maybe (StringTemplate a)
getStringTemplate' s sg = stGetFirst (sg s)

-- | Adds a set of global options to a single template
optInsertTmpl :: [(String, String)] -> StringTemplate a -> StringTemplate a
optInsertTmpl x st = st {senv = optInsert (map (second justSTR) x) (senv st)}

-- | Sets an encoding function of a template that all values are
-- rendered with. For example one useful encoder would be 'Text.Html.stringToHtmlString'. All attributes will be encoded once and only once.
setEncoder :: (Stringable a) => (String -> String) -> StringTemplate a -> StringTemplate a
setEncoder x st = st {senv = (senv st) {senc = x} }

-- | A special template that simply dumps the values of all the attributes set in it.
-- This may be made available to any template as a function by adding it to its group.
-- I.e. @ myNewGroup = addSuperGroup myGroup $ groupStringTemplates [("dumpAttribs", dumpAttribs)] @
dumpAttribs :: Stringable a => StringTemplate a
dumpAttribs = STMP (SEnv M.empty [] mempty id) $ \env -> showVal env (SM $ smp env)

{--------------------------------------------------------------------
  Internal API
--------------------------------------------------------------------}
--IMPLEMENT groups having stLookup return a Maybe for regions

data SEnv a = SEnv {smp :: SMap a, sopts :: [(String, (SEnv a -> SElem a))], sgen :: STGroup a, senc :: String -> String}

inSGen :: (STGroup a -> STGroup a) -> StringTemplate a -> StringTemplate a
inSGen f st@STMP{senv = env} = st {senv = env {sgen = f (sgen env)} }

envLookup :: String -> SEnv a -> Maybe (SElem a)
envLookup x = M.lookup x . smp
envInsert :: (String, SElem a) -> SEnv a -> SEnv a
envInsert (s, x) y = y {smp = M.insert s x (smp y)}
envInsApp :: Stringable a => String -> SElem a -> SEnv a -> SEnv a
envInsApp  s  x  y = y {smp = M.insertWith go s x (smp y)}
    where go a (LI bs) = LI (a:bs)
          go a b = LI [a,b]

optLookup :: String -> SEnv a -> Maybe (SEnv a -> SElem a)
optLookup x = lookup x . sopts
optInsert :: [(String, SEnv a -> SElem a)] -> SEnv a -> SEnv a
optInsert x env = env {sopts = x ++ sopts env}
nullOpt :: SEnv a -> SElem a
nullOpt = fromMaybe (justSTR "") =<< optLookup "null"

stLookup :: (Stringable a) => [Char] -> SEnv a -> StringTemplate a
stLookup x env = maybe (newSTMP ("No Template Found for: " ++ x))
                 (\st-> st {senv = mergeSEnvs env (senv st)}) $ stGetFirst (sgen env x)

--merges values of former into latter, preserving encoder
--of latter, as well as non-overriden options. group of latter is overridden.
mergeSEnvs :: SEnv a -> SEnv a -> SEnv a
mergeSEnvs x y = SEnv {smp = M.union (smp x) (smp y), sopts = (sopts y ++ sopts x), sgen = sgen x, senc = senc y}

parseSTMP :: (Stringable a) => (Char, Char) -> String -> SEnv a -> a
parseSTMP x = either (showStr .  show) (id) . runParser (stmpl False) (x,[]) ""

getSeps :: TmplParser (Char, Char)
getSeps = fst <$> getState

tellNames :: String -> TmplParser ()
tellNames x = getState >>= \(s,n) -> setState (s,x:n)

parseSTMPNames :: String -> Either ParseError [String]
parseSTMPNames = runParser getRefs (('$','$'),[]) ""
    where getRefs = do
            stmpl False :: TmplParser (SEnv String -> String)
            snd <$> getState

{--------------------------------------------------------------------
  Internal API for polymorphic display of elements
--------------------------------------------------------------------}

showVal :: Stringable a => SEnv a -> SElem a -> a
showVal snv se = case se of
                   STR x  -> stEncode x
                   BS  x  -> stFromByteString x
                   LI xs  -> joinUpWith showVal xs
                   SM sm  -> joinUpWith showAssoc $ M.assocs sm
                   STSH x -> stEncode (format x)
                   SBLE x -> x
                   SNull  -> showVal <*> nullOpt $ snv
    where sepVal = fromMaybe (justSTR "") =<< optLookup "separator" $ snv
          format = maybe stshow . stfshow <*> optLookup "format" $ snv
          joinUpWith f = mintercalate (showVal snv sepVal) . map (f snv)
          showAssoc e (k,v) = stEncode (k ++ ": ") `mlabel` showVal e v
          stEncode = stFromString . senc snv

showStr :: Stringable a => String -> SEnv a -> a
--showStr = flip showVal . STR
showStr s = const (stFromString $ s)

{--------------------------------------------------------------------
  Utility Combinators
--------------------------------------------------------------------}

justSTR :: String -> b -> SElem a
justSTR = const . STR
stshow :: STShow -> String
stshow (STShow a) = stringTemplateShow a
stfshow :: Stringable a => SEnv a -> (SEnv a -> SElem a) -> STShow -> String
stfshow snv fs (STShow a) = stringTemplateFormattedShow
                            (stToString <$$> showVal <*> fs $ snv) a

around :: Char -> GenParser Char st t -> Char -> GenParser Char st t
around x p y = do {char x; v<-p; char y; return v}
spaced :: GenParser Char st t -> GenParser Char st t
spaced p = do {spaces; v<-p; spaces; return v}
word :: GenParser Char st [Char]
word = many1 alphaNum
comlist :: GenParser Char st a -> GenParser Char st [a]
comlist p = spaced (p `sepBy1` spaced (char ','))

props :: Stringable a => TmplParser [SEnv a -> SElem a]
props = many $ char '.' >> (around '(' subexprn ')' <|> justSTR <$> word)

escapedChar, escapedStr :: [Char] -> GenParser Char st [Char]
escapedChar chs =
    noneOf chs >>= \x -> if x == '\\' then anyChar >>= \y ->
    if y `elem` chs then return [y] else return [x, y] else return [x]
escapedStr chs = concat <$> many1 (escapedChar chs)

{--------------------------------------------------------------------
  The Grammar
--------------------------------------------------------------------}

-- | if p is true, stmpl can fail gracefully, false it dies hard.
-- Set to false at the top level, and true within if expressions.
stmpl :: Stringable a => Bool -> TmplParser (SEnv a -> a)
stmpl p = do
  (ca, cb) <- getSeps
  mconcat <$> many (showStr <$> escapedStr [ca] <|> try (around ca optExpr cb)
                    <|> try comment <|> bl <?> "template")
      where bl | p = try blank | otherwise = blank

subStmp :: Stringable a => TmplParser (([SElem a], [SElem a]) -> SEnv a -> a)
subStmp = do
  (ca, cb) <- getSeps
  udEnv <- option (transform ["it"]) (transform <$> try attribNames)
  st <- mconcat <$> many (showStr <$> escapedStr (ca:"}|")
                         <|> try (around ca optExpr cb)
                         <|> try comment <|> blank  <?> "subtemplate")
  return (st <$$> udEnv)
      where transform an (att,is) =
                flip (foldr envInsert) $ zip ("i":"i0":an) (is++att)
            attribNames = (char '|' >>) . return =<< comlist (spaced word)

comment :: Stringable a => TmplParser (SEnv a -> a)
comment = do
  (ca, cb) <- getSeps
  string (ca:'!':[]) >> manyTill anyChar (try . string $ '!':cb:[])
  return (showStr "")

blank :: Stringable a => TmplParser (SEnv a -> a)
blank = do
  (ca, cb) <- getSeps
  char ca
  spaces
  char cb
  return (showStr "")

optExpr :: Stringable a => TmplParser (SEnv a -> a)
optExpr = do
  (_, cb) <- getSeps
  (try (string ("else"++[cb])) <|> try (string "elseif(") <|>
    try (string "endif")) .>> fail "Malformed If Statement." <|> return ()
  expr <- try ifstat <|> spaced exprn
  opts <- many opt
  skipMany (char ';')
  return $ expr . optInsert opts
      where opt = around ';' (spaced word) '=' >>= (<$> spaced subexprn) . (,)

{--------------------------------------------------------------------
  Statements
--------------------------------------------------------------------}

getProp :: Stringable a => [SEnv a -> SElem a] -> SElem a -> SEnv a -> SElem a
getProp (p:ps) (SM mp) env =
  case M.lookup (stToString . showVal env $ p env) mp of
    Just prop -> getProp ps prop env
    Nothing -> SNull
getProp (_:_) _ _ = SNull
getProp _ se _ = se

ifIsSet :: t -> t -> Bool -> SElem a -> t
ifIsSet t e n SNull = if n then e else t
ifIsSet t e n _ = if n then t else e

ifstat ::Stringable a => TmplParser (SEnv a -> a)
ifstat = do
  (_, cb) <- getSeps
  string "if("
  n <- option True (char '!' >> return False)
  e <- subexprn
  p <- props
  char ')' >> char cb
  act <- stmpl True
  cont <- (try elseifstat <|> try elsestat <|> endifstat)
  return (ifIsSet act cont n =<< getProp p =<< e)

elseifstat ::Stringable a => TmplParser (SEnv a -> a)
elseifstat = getSeps >>= char . fst >> string "else" >> ifstat

elsestat ::Stringable a => TmplParser (SEnv a -> a)
elsestat = do
  (ca, cb) <- getSeps
  around ca (string "else") cb
  act <- stmpl True
  char ca >> string "endif"
  return act

endifstat ::Stringable a => TmplParser (SEnv a -> a)
endifstat = getSeps >>= char . fst >> string "endif" >> return (showStr "")

{--------------------------------------------------------------------
  Expressions
--------------------------------------------------------------------}

exprn :: Stringable a => TmplParser (SEnv a -> a)
exprn = do
  exprs <- comlist subexprn <?> "expression"
  templ <- many (char ':' >> iterApp <$> comlist (anonTmpl <|> regTemplate)) <?> "template call"
  return $ fromMany (showVal <*> head exprs)
             ((sequence exprs >>=) . seqTmpls) templ

seqTmpls :: Stringable a => [[SElem a] -> SEnv a -> a] -> [SElem a] -> SEnv a -> a
seqTmpls [f]    y = f y
seqTmpls (f:fs) y = seqTmpls fs =<< (:[]) . SBLE . f y
seqTmpls  _ _     = const (stFromString "")

subexprn :: Stringable a => TmplParser (SEnv a -> SElem a)
subexprn = cct <$> spaced
            (braceConcat
             <|> SBLE <$$> ($ ([SNull],ix0)) <$> try regTemplate
             <|> attrib
             <|> SBLE <$$> ($ ([SNull],ix0)) <$> anonTmpl
             <?> "expression")
           `sepBy1` spaced (char '+')
    where cct xs@(_:_:_) = SBLE |.
                           flip mconcatMap <$> showVal <*> sequence xs
          cct [x] = x
          cct  _  = const SNull

braceConcat :: Stringable a => TmplParser (SEnv a -> SElem a)
braceConcat = LI . foldr go [] <$$> sequence <$> around '['(comlist subexprn)']'
    where go (LI x) lst = x++lst; go x lst = x:lst

literal :: GenParser Char st (b -> SElem a)
literal = justSTR <$> (around '"' (concat <$> many (escapedChar "\"")) '"'
                   <|> around '\'' (concat <$> many (escapedChar "'")) '\'')

attrib :: Stringable a => TmplParser (SEnv a -> SElem a)
attrib = do
  a <-     literal
       <|> try functn
       <|> prepExp <$> word
       <|> prepExp <$> qqWord
       <|> around '(' subexprn ')'
          <?> "attribute"
  proprs <- props
  return $ fromMany a ((a >>=) . getProp) proprs
      where prepExp var = fromMaybe SNull <$> envLookup var
            qqWord = do
              w <- around '`' word '`'
              tellNames w
              return $ '`' : w ++ "`"

--add null func
functn :: Stringable a => TmplParser (SEnv a -> SElem a)
functn = do
  f <- string "first" <|> string "rest" <|> string "strip"
       <|> try (string "length") <|> string "last" <?> "function"
  (fApply f .) <$> around '(' subexprn ')'
      where fApply str (LI xs)
                | str == "first"  = if null xs then SNull else head xs
                | str == "last"   = last xs
                | str == "rest"   = if null xs then SNull else (LI . tail) xs
                | str == "strip"  = LI . filter (not . liNil) $ xs
                | str == "length" = STR . show . length $ xs
            fApply str x
                | str == "rest"   = (LI [])
                | str == "length" = STR "1"
                | otherwise       = x
            liNil (LI x) = null x
            liNil _      = False

{--------------------------------------------------------------------
  Templates
--------------------------------------------------------------------}
--change makeTmpl to do notation for clarity?

mkIndex :: Num b => [b] -> [[SElem a]]
mkIndex = map ((:) . STR . show . (1+) <*> (:[]) . STR . show)
ix0 :: [SElem a]
ix0 = [STR "1",STR "0"]

cycleApp :: (Stringable a) => [([SElem a], [SElem a]) -> SEnv a -> a] -> [([SElem a], [SElem a])]  -> SEnv a -> a
cycleApp x y = mconcatMap (zipWith ($) (cycle x) y) . flip ($)

pluslen :: [a] -> [([a], [SElem b])]
pluslen xs = zip (map (:[]) xs) $ mkIndex [0..(length xs)]

liTrans :: [SElem a] -> [([SElem a], [SElem a])]
liTrans = pluslen' . paddedTrans SNull . map u
    where u (LI x) = x; u x = [x]
          pluslen' xs = zip xs $ mkIndex [0..(length xs)]

iterApp :: Stringable a => [([SElem a], [SElem a]) -> SEnv a -> a] -> [SElem a] -> SEnv a -> a
iterApp [f] (LI xs:[])    = (mconcatMap $ pluslen xs) . flip f
iterApp [f] vars@(LI _:_) = (mconcatMap $ liTrans vars) . flip f
iterApp [f] v             = f (v,ix0)
iterApp fs (LI xs:[])     = cycleApp fs (pluslen xs)
iterApp fs vars@(LI _:_)  = cycleApp fs (liTrans vars)
iterApp fs xs             = cycleApp fs (pluslen xs)

anonTmpl :: Stringable a => TmplParser (([SElem a], [SElem a]) -> SEnv a -> a)
anonTmpl = around '{' subStmp '}'

regTemplate :: Stringable a => TmplParser (([SElem a], [SElem a]) -> SEnv a -> a)
regTemplate = do
  try (functn::TmplParser (SEnv String -> SElem String)) .>> fail "" <|> return ()
  name <- justSTR <$> many1 (alphaNum <|> char '/'<|> char '_')
          <|> around '(' subexprn ')'
  vals <- around '(' (spaced $ try assgn <|> anonassgn <|> return []) ')'
  return $ join . (. name) . makeTmpl vals
      where makeTmpl v ((se:_),is) (STR x)  =
                render |. stBind . (zip ["it","i","i0"] (se:is) ++)
                           . swing (map . second) v <*> stLookup x
            makeTmpl _ _ _ = showStr "Invalid Template Specified"
            stBind v st = st {senv = foldr envInsert (senv st) v}
            anonassgn = (:[]) . (,) "it" <$> subexprn
            assgn = (spaced word >>= (<$> char '=' .>> spaced subexprn) . (,))
                    `sepEndBy1` char ';'

--DEBUG

{-pTrace s = pt <|> return ()
    where pt = try $
               do
                 x <- try $ many1 anyChar
                 trace (s++": " ++x) $ try $ char 'z'
                 fail x
-}