{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
module Pages.TH where

import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Char (toUpper)
import Data.List.NonEmpty (NonEmpty(..))
import Data.List.NonEmpty qualified as NE
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Text.Builder.Linear.Buffer
import Language.Haskell.TH
import Text.Mustache qualified as M
import Text.Mustache.Types qualified as M

import EscapeXML


data DataMap
  = VQInt
  | VQBool
  | VQText
  | VQUnit  -- ()
  | VQList DataMap  -- ^ embedded map is of a list element
  | VQMaybe DataMap
  | VQObject [(Text, (Exp -> Exp, DataMap))]  -- ^ field name, projection function, submap


makeDataMap :: Name -> Q DataMap
makeDataMap name = makeDataMapTy (ConT name)

makeDataMapTy :: Type -> Q DataMap
makeDataMapTy (ConT tyname)
  | tyname == ''Int = return VQInt
  | tyname == ''Bool = return VQBool
  | tyname == ''Text = return VQText
makeDataMapTy (TupleT 0) = return VQUnit
makeDataMapTy (AppT ListT ty) = VQList <$> makeDataMapTy ty
makeDataMapTy (AppT (ConT name) ty) | name == ''Maybe = VQMaybe <$> makeDataMapTy ty

makeDataMapTy topty = do
  case splitApps topty [] of
    (ConT tyname, tyargs) ->
      reify tyname >>= \case
        TyConI (DataD [] _ tyvars _ [RecC conname fields] _) ->
          let n = length fields
          in VQObject <$> sequence
               [do dm <- makeDataMapTy (subst (Map.fromList (zip (map tyvarName tyvars) tyargs)) fieldty)
                   f <- projectField conname i n
                   return (T.pack (nameBase name), (f, dm))
               | ((name, _bang, fieldty), i) <- zip fields [0..]]
        _ -> fail $ "Not a suitable data type: " ++ show tyname
    _ -> fail $ "Not a straightforward type: " ++ show topty
  where
    projectField :: Name -> Int -> Int -> Q (Exp -> Exp)
    projectField conname i n = do
      varname <- newName ('x' : show (i + 1))
      let pats = replicate i WildP ++ [VarP varname] ++ replicate (n-i-1) WildP
      return $ \e ->
        CaseE e [Match (ConP conname [] pats) (NormalB (VarE varname)) []]

    splitApps :: Type -> [Type] -> (Type, [Type])
    splitApps (AppT t1 t2) args = splitApps t1 (t2 : args)
    splitApps t args = (t, args)

    tyvarName :: TyVarBndr flag -> Name
    tyvarName (PlainTV n _) = n
    tyvarName (KindedTV n _ _) = n

    subst :: Map Name Type -> Type -> Type
    subst _ t@ConT{} = t
    subst _ t@ListT = t
    subst _ t@TupleT{} = t
    subst mp (VarT n) | Just t <- Map.lookup n mp = t
    subst mp (AppT t1 t2) = AppT (subst mp t1) (subst mp t2)
    subst _ t = error $ "Unsupported type: " ++ show t

makeRender :: M.Template -> Name -> Q [Dec]
makeRender (M.Template tplname nodes _) dataname = do
  dm <- makeDataMap dataname
  let funname = mkName ("renderPage" ++ toPascalCase tplname)
  argname <- newName "dat"
  bufname <- newName "buf"
  body <- renderNodes ((dm, VarE argname) :| []) (DuplicableExp (VarE bufname)) nodes
  let body' = VarE 'runBufferBS `AppE` LamE [VarP bufname] body
  return
    [SigD funname (ArrowT `AppT` ConT dataname `AppT` ConT ''ByteString)
    ,FunD funname
      [Clause [BangP (VarP argname)] (NormalB body') []]]

data ExpensiveExp
  = ExpensiveExp Exp
  | DuplicableExp Exp

expensive :: ExpensiveExp -> Exp
expensive (ExpensiveExp e) = e
expensive (DuplicableExp e) = e

makeDuplicable :: ExpensiveExp -> (Exp -> Q Exp) -> Q Exp
makeDuplicable (ExpensiveExp e) k = do
  varname <- newName "a"
  caseE (pure e)
    [match (varP varname) (normalB (k (VarE varname))) []]
makeDuplicable (DuplicableExp e) k = k e

-- | This is just '(|>)', but with a @NOINLINE@ pragma. This is because GHC
-- eagerly inlines '(|>)', and that function contains some conditionals, and as
-- a result the generated Core for the render function contains a whole forest
-- of join points. By avoiding inlining this function, the Core becomes a
-- straight line of @case@s that contains join points only when there is
-- genuinely control flow going on, i.e. when substituting a @Maybe@ or when
-- looping over a list.
{-# NOINLINE noinlineAppText #-}
noinlineAppText :: Buffer %1-> Text -> Buffer
noinlineAppText = (|>)

-- | All input expressions must be duplicable.
renderNodes :: NonEmpty (DataMap, Exp) -> ExpensiveExp -> [M.Node Text] -> Q (Exp {- Buffer -})
renderNodes _    buf [] = return (expensive buf)
renderNodes dats buf (node : nodes) = do
  let app buf' name e = InfixE (Just buf') (VarE name) (Just e)

  buf' <- case node of
    M.TextBlock t ->
      let bytes = BS.unpack (TE.encodeUtf8 t)
      in return $ app (expensive buf) '(|>#) (LitE (StringPrimL bytes))
    M.Variable esc named -> 
      case search named dats of
        Just (VQInt, e) -> return $ app (expensive buf) '(|>$) e
        Just (VQText, e) ->
          return $ app (expensive buf) 'noinlineAppText (if esc then VarE 'escapeXML `AppE` e else e)
        Just (_, _) -> fail $ "Strange value in variable reference " ++ show named
        Nothing -> fail $ "Variable not found: " ++ show named
    M.Section named nodes' ->
      case search named dats of
        Just (VQBool, e) ->
          makeDuplicable buf $ \buf' ->
            condE (pure e)
              (renderNodes ((VQBool, e) NE.<| dats) (DuplicableExp buf') nodes')
              (pure buf')
        Just (VQText, e) ->
          makeDuplicable buf $ \buf' ->
            condE [| T.null $(pure e) |]
              (pure buf')
              (renderNodes ((VQText, e) NE.<| dats) (DuplicableExp buf') nodes')
        Just (VQList dm, e) -> do
          bvar <- newName "b"
          vvar <- newName "v"
          varE 'foldlIntoBuffer
            `appE` lamE [varP bvar, varP vvar]
                     (renderNodes ((dm, VarE vvar) NE.<| dats) (DuplicableExp (VarE bvar)) nodes')
            `appE` pure (expensive buf)
            `appE` pure e
        Just (VQMaybe dm, e) -> do
          xvar <- newName "x"
          makeDuplicable buf $ \buf' ->
            caseE (pure e)
              [match (conP 'Nothing []) (normalB (pure buf')) []
              ,match (conP 'Just [varP xvar])
                     (normalB (renderNodes ((dm, VarE xvar) NE.<| dats) (DuplicableExp buf') nodes'))
                     []]
        Just pair@(VQObject _, _) -> do
          renderNodes (pair NE.<| dats) buf nodes'
        Just (VQInt, _) -> fail $ "Int value in section scrutinee " ++ show named
        Just (VQUnit, _) -> fail $ "() value in section scrutinee " ++ show named
        Nothing -> fail $ "Section scrutinee not found: " ++ show named
    M.InvertedSection named nodes' ->
      case search named dats of
        Just (VQBool, e) ->
          makeDuplicable buf $ \buf' ->
            condE (pure e)
              (pure buf')
              (renderNodes dats (DuplicableExp buf') nodes')
        Just (VQText, e) ->
          makeDuplicable buf $ \buf' ->
            condE [| T.null $(pure e) |]
              (renderNodes dats (DuplicableExp buf') nodes')
              (pure buf')
        Just (VQMaybe _, e) ->
          makeDuplicable buf $ \buf' ->
            caseE (pure e)
              [match (conP 'Nothing []) (normalB (renderNodes dats (DuplicableExp buf') nodes')) []
              ,match (conP 'Just [wildP]) (normalB (pure buf')) []]
        Just (VQList _, e) ->
          makeDuplicable buf $ \buf' ->
            caseE (pure e)
              [match (listP []) (normalB (renderNodes dats (DuplicableExp buf') nodes')) []
              ,match wildP (normalB (pure buf')) []]
        Just (VQObject _, _) -> fail $ "Object value in inverted section scrutinee " ++ show named
        Just (VQInt, _) -> fail $ "Int value in inverted section scrutinee " ++ show named
        Just (VQUnit, _) -> fail $ "() value in inverted section scrutinee " ++ show named
        Nothing -> fail $ "Inverted section scrutinee not found: " ++ show named
    M.Partial _ _ -> fail $ "TODO partials"

  renderNodes dats (ExpensiveExp buf') nodes

search :: M.DataIdentifier -> NonEmpty (DataMap, Exp) -> Maybe (DataMap, Exp)
search M.Implicit (value :| _) = Just value
search (M.NamedData []) (value :| _) = Just value
search (M.NamedData (key : keys)) ((dm, e) :| dats) =
  case dm of
    VQObject fields ->
      case lookup key fields of
        Just (proj, dm') -> search (M.NamedData keys) ((dm', proj e) :| [])
        Nothing -> search (M.NamedData (key : keys)) =<< NE.nonEmpty dats
    _ ->
      search (M.NamedData (key : keys)) =<< NE.nonEmpty dats

toPascalCase :: String -> String
toPascalCase "" = ""
toPascalCase (c:cs) = toUpper c : go cs
  where go ('-':cs') = toPascalCase cs'
        go (c':cs') = c' : go cs'
        go "" = ""