{-# LANGUAGE TemplateHaskellQuotes #-}
module Simplify.TH (simprec) where

import Data.Bifunctor (first)
import Data.Char
import Data.List (foldl1')
import Language.Haskell.TH
import Language.Haskell.TH.Quote
import Text.ParserCombinators.ReadP


-- [simprec| EPair ext *a *b |]
-- ~>
-- do a' <- within (\a' -> EPair ext a' b) (simplify' a)
--    b' <- within (\b' -> EPair ext a' b') (simplify' b)
--    pure (EPair ext a' b')

simprec :: QuasiQuoter
simprec = QuasiQuoter
  { quoteDec = \_ -> fail "simprec used outside of expression context"
  , quoteType = \_ -> fail "simprec used outside of expression context"
  , quoteExp = handler
  , quotePat = \_ -> fail "simprec used outside of expression context"
  }

handler :: String -> Q Exp
handler str =
  case readP_to_S pTemplate str of
    [(template, "")] -> generate template
    _:_:_ -> fail "simprec: template grammar ambiguous"
    _ -> fail "simprec: could not parse template"

generate :: Template -> Q Exp
generate (Template topitems) =
  let takePrefix (Plain x : xs) = first (x:) (takePrefix xs)
      takePrefix xs = ([], xs)

      itemVar "" = error "simprec: empty item name?"
      itemVar name@(c:_) | isLower c = VarE (mkName name)
                         | isUpper c = ConE (mkName name)
                         | otherwise = error "simprec: non-letter item name?"

      loop :: Exp -> [Item] -> Q [Stmt]
      loop yet [] = return [NoBindS (VarE 'pure `AppE` yet)]
      loop yet (Plain x : xs) = loop (yet `AppE` itemVar x) xs
      loop yet (Recurse x : xs) = do
        primeName <- newName (x ++ "'")
        let appPrePrime e (Plain y) = e `AppE` itemVar y
            appPrePrime e (Recurse y) = e `AppE` itemVar y
        let stmt = BindS (VarP primeName) $
                     VarE (mkName "within")
                     `AppE` LamE [VarP primeName] (foldl' appPrePrime (yet `AppE` VarE primeName) xs)
                     `AppE` (VarE (mkName "simplify'") `AppE` VarE (mkName x))
        stmts <- loop (yet `AppE` VarE primeName) xs
        return (stmt : stmts)

      (prefix, items') = takePrefix topitems
  in DoE Nothing <$> loop (foldl1' AppE (map itemVar prefix)) items'

data Template = Template [Item]
  deriving (Show)

data Item = Plain String | Recurse String
  deriving (Show)

pTemplate :: ReadP Template
pTemplate = do
  items <- many (skipSpaces >> pItem)
  skipSpaces
  eof
  return (Template items)

pItem :: ReadP Item
pItem = (char '*' >> Recurse <$> pName) +++ (Plain <$> pName)

pName :: ReadP String
pName = do
  c1 <- satisfy (\c -> isAlpha c || c == '_')
  cs <- munch (\c -> isAlphaNum c || c `elem` "_'")
  return (c1:cs)