{-# LANGUAGE TemplateHaskell #-}
module BuilderQQ
  (b,
   Builder, toLazyText, fromText, fromString, fromChar,
   textMap, ifJust)
where

import Data.Char (isLower, isSpace, isDigit, isAlphaNum)
import Language.Haskell.TH
import Language.Haskell.TH.Quote
import Data.Maybe (mapMaybe)
import Data.Text.Lazy.Builder
  (Builder, fromText, fromString, singleton, toLazyText)
import Text.Read (readMaybe)
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Text.Lazy (toStrict)
import Data.Foldable
import Data.Semigroup

data ChunkType = Lit | Var VarType deriving Show
data VarType =
    Plain
  | FromText
  | FromString
  | FromChar
  | Show
  | Reindent !Int
  | ReindentList !Int
  deriving Show
type Chunk = (ChunkType, Text)

indent :: Int -> Text -> Builder
indent i str
  | Text.all isSpace str = ""
  | otherwise            = replicateB i ' ' <> fromText str

reindent :: Int -> Text -> Builder
reindent i str =
  fold $ mapInit (<> "\n") $
  map2 (fromText . dropIndent) (indent i . dropIndent) ls
 where
  ls  = dropWhile (Text.all isSpace) $ Text.lines str
  ls' = filter (Text.any $ not . isSpace) ls

  dropIndent = Text.drop minIndent

  minIndent =
    getMin $ option 0 id $ foldMap (Option . Just . Min . indentOf) ls'

  indentOf = go 0 where
    go n (' '  :. cs) = go (n + 1) cs
    go n ('\t' :. cs) = go (((n `mod` 8) + 1) * 8) cs
    go n _            = n

  map2 _ _ []     = []
  map2 f g (x:xs) = f x : map g xs


chunks :: Text -> [Chunk]
chunks = reverse . go "" [] . trimEnd where
  go acc cs NilT = if Text.null acc' then cs else (Lit, acc') : cs
    where acc' = toStrictText acc

  -- $&: expands to nothing
  go acc cs ('$' :. '&' :. rest) = go acc cs rest

  -- $$: expands to one $
  go acc cs ('$' :. '$' :. rest) = go (acc <> "$") cs rest

  -- $*var: expands to (fromText $var)
  go acc cs ('$' :. '*' :. rest) =
      go "" ((Var FromText, var) : lit acc : cs) rest2
    where (var, rest2) = splitVar rest

  -- $@var: expands to (fromString $var)
  go acc cs ('$' :. '@' :. rest) =
      go "" ((Var FromString, var) : lit acc : cs) rest2
    where (var, rest2) = splitVar rest

  -- $'var: expands to (singleton $var)
  go acc cs ('$' :. '\'' :. rest) =
      go "" ((Var FromChar, var) : lit acc : cs) rest2
    where (var, rest2) = splitVar rest

  -- $^var: expands to (fromString (show $var))
  go acc cs ('$' :. '^' :. rest) =
      go "" ((Var Show, var) : lit acc : cs) rest2
    where (var, rest2) = splitVar rest

  -- $n*var (n a number): expands to builder var indented by n
  -- $n.var: same but var is a list
  go acc cs ('$' :. rest@(d :. _)) | isDigit d =
      go "" ((Var ty, var) : lit acc : cs) rest3
    where
      (n', c :. rest2) = Text.span isDigit rest
      n = read $ Text.unpack n'
      (var, rest3) = splitVar rest2
      ty = case c of
        '*' -> Reindent n
        '.' -> ReindentList n
        _   -> error $ "unknown reindent type " ++ show c

  -- $var: expands to that var's contents
  go acc cs ('$' :. rest) =
      go "" ((Var Plain, var) : lit acc : cs) rest2
    where (var, rest2) = splitVar rest

  go acc cs (c :. rest) = go (acc <> singleton c) cs rest

  splitVar s
    | (var@(v :. _), s') <- Text.span isIdChar s,
      isLower v || v == '_'
    = (var, s')
  splitVar _ = error "invalid variable name"

  isIdChar c = isAlphaNum c || c `elem` ("_'" :: String)

  lit s = (Lit, toStrictText s)

  trimEnd = Text.dropWhileEnd isSpace

chunksWithReindent :: String -> [Chunk]
chunksWithReindent ('@':str)
  | [(n',rest)] <- lex str,
    Just n      <- readMaybe n'
  = chunks $ toStrictText $ reindent n $ Text.pack rest
chunksWithReindent str = chunks $ Text.pack str

toStrictText :: Builder -> Text
toStrictText = toStrict . toLazyText


chunksToExpQ :: [Chunk] -> ExpQ
chunksToExpQ cs = [|mconcat $es|] where
  es = listE $ mapMaybe chunk1 cs
  chunk1 (Lit, "")     = Nothing
  chunk1 (Lit, lit)    = Just $ stringE $ Text.unpack lit
  chunk1 (Var t, name) = Just $ case t of
      Plain          -> var
      FromText       -> [|fromText $var|]
      FromString     -> [|fromString $var|]
      FromChar       -> [|singleton $var|]
      Show           -> [|fromString $ show $var|]
      Reindent n     -> [|reindent n $var|]
      ReindentList n -> [|reindentList n $var|]
    where var = varE (mkName $ Text.unpack name)

reindentList :: Int -> [Builder] -> Builder
reindentList n = fold . mapInit (<> "\n") . mapTail (replicateB n ' ' <>)

replicateB :: Int -> Char -> Builder
replicateB n c = fromText $ Text.replicate n $ Text.singleton c

mapInit :: (a -> a) -> [a] -> [a]
mapInit _ []  = []
mapInit _ [x] = [x]
mapInit f (x:xs) = f x : mapInit f xs

mapTail :: (a -> a) -> [a] -> [a]
mapTail _ [] = []
mapTail f (x:xs) = x : map f xs

b :: QuasiQuoter
b = QuasiQuoter {
    quoteExp  = chunksToExpQ . chunksWithReindent,
    quotePat  = error "can't use in a pattern",
    quoteType = error "can't use in a type",
    quoteDec  = error "can't use at top level"
  }


pattern NilT :: Text
pattern NilT <- (Text.null -> True)
  where NilT = ""

infixr 5 :.
pattern (:.) :: Char -> Text -> Text
pattern c :. t <- (Text.uncons -> Just (c, t))
  where c :. t = Text.cons c t

{-# COMPLETE NilT, (:.) :: Text #-}

fromChar :: Char -> Builder
fromChar = singleton

textMap :: (Char -> Builder) -> Text -> Builder
textMap f = Text.foldl' (\buf c -> buf <> f c) mempty

ifJust :: Monoid b => Maybe a -> (a -> b) -> b
ifJust x f = maybe mempty f x