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

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

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

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

reindentB :: Int -> Builder -> Builder
reindentB i (toLazyText -> str) =
  fold $ intersperse "\n" $
  map2 (fromLazyText . dropIndent) (indent i . dropIndent) ls
 where
  ls  = dropWhile (LText.all isSpace) $ LText.lines str
  ls' = filter (LText.any $ not . isSpace) ls

  dropIndent = LText.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

  -- $n.var (n a number): expands to builder var indented by n
  go acc cs ('$' :. rest@(d :. _)) | isDigit d =
      go "" ((Var (Reindent n), var) : lit acc : cs) rest3
    where
      ((read . Text.unpack -> n), '.' :. rest2) = Text.span isDigit rest
      (var, rest3) = splitVar rest2

  -- $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 :: Builder|] 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          -> [|build $var|]
      Reindent n     -> [|reindent n $var|]
    where var = varE (mkName $ Text.unpack name)

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

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 #-}

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


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


escId :: Text -> Builder
escId = foldMap esc1 . Text.unpack where
  esc1 c | isSpace c       = ""
         | latin1Special c = "_"
         | otherwise       = fromChar c
  latin1Special c =
    c <= 'ÿ' && not (isAlphaNum c) && c /= '-'

escAttr :: Text -> Builder
escAttr = foldMap esc1 . Text.unpack where
  esc1 '<'  = "&lt;"
  esc1 '>'  = "&gt;"
  esc1 '"'  = "&quot;"
  esc1 '\'' = "&apos;"
  esc1 c    = fromChar c


class CanBuild a where
  build :: a -> Builder
  reindent :: Int -> a -> Builder
  reindent i = reindentB i . build

instance CanBuild Builder    where build = id
instance CanBuild Text       where build = fromText
instance CanBuild LText.Text where build = fromLazyText
instance CanBuild Char       where build = singleton
instance CanBuild String     where build = fromString

newtype ShowBuild a = ShowBuild a deriving newtype Show
instance Show a => CanBuild (ShowBuild a) where build = build . show

deriving via ShowBuild Int     instance CanBuild Int
deriving via ShowBuild Integer instance CanBuild Integer

instance {-# OVERLAPPABLE #-} CanBuild a => CanBuild [a] where
  build = foldMap \x -> build x <> "\n"
  reindent n = fold . intersperse ("\n" <> replicateB n ' ') . map build