{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}

module Jikka.Core.Language.QuasiRules
  ( r,

    -- * Things which `r` uses.
    module Jikka.Core.Language.Expr,
    alphaExpr,
    makeRewriteRule,
    genVarName',
  )
where

import Control.Arrow
import Control.Monad.State.Strict
import Data.Data
import Data.Maybe
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.Common.Format.Error
import qualified Jikka.Core.Convert.Alpha as Alpha
import qualified Jikka.Core.Convert.TypeInfer as TypeInfer
import Jikka.Core.Language.Expr
import Jikka.Core.Language.RewriteRules
import Jikka.Core.Language.Util
import Jikka.Core.Parse (parseRule)
import Language.Haskell.TH (Exp (..), Lit (..), Pat (..), Q, Stmt (..))
import qualified Language.Haskell.TH as TH
import qualified Language.Haskell.TH.Quote as TH
import qualified Language.Haskell.TH.Syntax as TH

liftError :: ExceptT Error Q a -> Q a
liftError f = do
  x <- runExceptT f
  case x of
    Left err -> fail $ "Jikka.Core.Language.QuasiRules.liftError: " ++ unlines (prettyError' err)
    Right y -> return y

fromVarName :: VarName -> Q TH.Name
fromVarName (VarName x _) = do
  TH.newName (fromMaybe "x" x)

fromTypeName :: TypeName -> Q TH.Name
fromTypeName (TypeName x _) = do
  TH.newName (fromMaybe "t" x)

liftDataP :: Data a => a -> Q Pat
liftDataP = TH.dataToPatQ (const Nothing)

-- | `Exp` with type `Expr`.
type ExprExp = Exp

data RenamedVarName
  = DeclaredAtForall
  | RenamedVar ExprExp
  | RenamedPatLam TH.Name ExprExp
  | RenamedExpLam ExprExp
  | RenamedPatLet TH.Name ExprExp
  | RenamedExpLet ExprExp
  deriving (Eq, Ord, Show)

expFromRenamedVarName :: RenamedVarName -> Maybe ExprExp
expFromRenamedVarName = \case
  DeclaredAtForall -> Nothing
  RenamedVar e -> Just e
  RenamedPatLam _ e -> Just e
  RenamedExpLam e -> Just e
  RenamedPatLet _ e -> Just e
  RenamedExpLet e -> Just e

data Env = Env
  { vars :: [(VarName, RenamedVarName)],
    typeVars :: [(TypeName, TH.Name)]
  }

toPatT :: Type -> StateT Env Q Pat
toPatT = \case
  VarTy x -> do
    env <- gets typeVars
    case lookup x env of
      Just y -> do
        lift [p|((==) $(pure (VarE y)) -> True)|]
      Nothing -> do
        y <- lift $ fromTypeName x
        modify' (\env -> env {typeVars = (x, y) : typeVars env})
        return $ VarP y
  IntTy -> lift $ liftDataP IntTy
  BoolTy -> lift $ liftDataP IntTy
  ListTy t -> do
    t <- toPatT t
    lift [p|ListTy $(pure t)|]
  TupleTy ts -> do
    ts <- mapM toPatT ts
    lift [p|TupleTy $(pure (ListP ts))|]
  FunTy t1 t2 -> do
    t1 <- toPatT t1
    t2 <- toPatT t2
    lift [p|FunTy $(pure t1) $(pure t2)|]
  DataStructureTy ds -> do
    lift [p|DataStructureTy $(liftDataP ds)|]

toPatL :: Literal -> StateT Env Q Pat
toPatL = \case
  LitBuiltin builtin ts -> do
    ts <- mapM toPatT ts
    lift [p|LitBuiltin $(liftDataP builtin) $(pure (ListP ts))|]
  lit@(LitInt _) -> lift $ liftDataP lit
  lit@(LitBool _) -> lift $ liftDataP lit
  LitNil t -> do
    t <- toPatT t
    lift [p|LitNil $(pure t)|]
  LitBottom t msg -> do
    t <- toPatT t
    lift [p|LitBottom $(pure t) $(liftDataP msg)|]

toPatE :: Expr -> StateT Env Q Pat
toPatE = \case
  Var x ->
    if x == VarName Nothing Nothing
      then return WildP
      else do
        env <- gets vars
        case expFromRenamedVarName <$> lookup x env of
          Just (Just y) -> do
            lift [p|((== $(pure y)) -> True)|]
          Just Nothing -> do
            y <- lift $ fromVarName x
            modify' (\env -> env {vars = (x, RenamedVar (VarE y)) : vars env})
            return $ VarP y
          Nothing -> fail $ "Jikka.Core.Language.QuasiRules.toPatE: undefined variable (forall is required): " ++ formatVarName x
  Lit lit -> do
    lit <- toPatL lit
    lift [p|Lit $(pure lit)|]
  App e1 e2 -> do
    e1 <- toPatE e1
    e2 <- toPatE e2
    lift [p|App $(pure e1) $(pure e2)|]
  Lam x t e -> do
    t <- toPatT t
    y <- lift $ fromVarName x
    y' <- lift [e|Var $(pure (VarE y))|]
    modify' (\env -> env {vars = (x, RenamedPatLam y y') : vars env})
    e <- toPatE e
    lift [p|Lam $(pure (VarP y)) $(pure t) $(pure e)|]
  Let x t e1 e2 -> do
    t <- toPatT t
    e1 <- toPatE e1
    y <- lift $ fromVarName x
    y' <- lift [e|Var $(pure (VarE y))|]
    modify' (\env -> env {vars = (x, RenamedPatLet y y') : vars env})
    e2 <- toPatE e2
    lift [p|Let $(pure (VarP y)) $(pure t) $(pure e1) $(pure e2)|]
  Assert e1 e2 -> do
    e1 <- toPatE e1
    e2 <- toPatE e2
    lift [p|Assert $(pure e1) $(pure e2)|]

toExpT :: Type -> StateT Env Q Exp
toExpT = \case
  VarTy x -> do
    env <- gets typeVars
    case lookup x env of
      Just y -> return $ VarE y
      Nothing -> fail $ "Jikka.Core.Language.QuasiRules.toExpT: undefined type variable: " ++ formatTypeName x
  IntTy -> do
    lift $ TH.liftData IntTy
  BoolTy -> do
    lift $ TH.liftData BoolTy
  ListTy t -> do
    t <- toExpT t
    lift [e|ListTy $(pure t)|]
  TupleTy ts -> do
    ts <- mapM toExpT ts
    lift [e|TupleTy $(pure (ListE ts))|]
  FunTy t1 t2 -> do
    t1 <- toExpT t1
    t2 <- toExpT t2
    lift [e|FunTy $(pure t1) $(pure t2)|]
  DataStructureTy ds -> do
    lift $ TH.liftData (DataStructureTy ds)

toExpL :: Literal -> StateT Env Q Exp
toExpL = \case
  LitBuiltin builtin ts -> do
    ts <- mapM toExpT ts
    lift [e|LitBuiltin $(TH.liftData builtin) $(pure (ListE ts))|]
  lit@(LitInt _) -> lift $ TH.liftData lit
  lit@(LitBool _) -> lift $ TH.liftData lit
  LitNil t -> do
    t <- toExpT t
    lift [e|LitNil $(pure t)|]
  LitBottom t msg -> do
    t <- toExpT t
    lift [e|LitBottom $(pure t) $(TH.liftData msg)|]

toExpE :: Expr -> StateT Env Q ([Stmt], Exp)
toExpE e = do
  var <- lift [e|Var|]
  genVarName <- lift [e|genVarName'|]
  case e of
    Var x -> do
      env <- gets vars
      case expFromRenamedVarName <$> lookup x env of
        Just (Just y) -> return ([], y)
        _ -> fail $ "Jikka.Core.Language.QuasiRules.toExpE: undefined variable: " ++ formatVarName x
    Lit lit -> do
      lit <- toExpL lit
      e <- lift [e|Lit $(pure lit)|]
      return ([], e)
    App e1 e2 -> do
      (stmts, e1) <- toExpE e1
      (stmts', e2) <- toExpE e2
      e <- lift [e|App $(pure e1) $(pure e2)|]
      return (stmts ++ stmts', e)
    Lam x t e -> do
      t <- toExpT t
      y <- gets (lookup x . vars)
      case y of
        Just (RenamedPatLam y _) -> do
          -- Use the same variable name
          (stmts, e) <- toExpE e
          e <- lift [e|Lam $(pure (VarE y)) $(pure t) $(pure e)|]
          return (stmts, e)
        Nothing -> do
          -- Introduce a new name
          y <- lift $ fromVarName x
          modify' (\env -> env {vars = (x, RenamedExpLam (AppE var (VarE y))) : vars env})
          (stmts, e) <- toExpE e
          e <- lift [e|Lam $(pure (VarE y)) $(pure t) $(pure e)|]
          return (BindS (VarP y) genVarName : stmts, e)
        _ -> fail $ "Jikka.Core.Language.QuasiRules.toExpE: variable conflicts: " ++ formatVarName x
    Let x t e1 e2 -> do
      t <- toExpT t
      (stmts, e1) <- toExpE e1
      y <- gets (lookup x . vars)
      case y of
        Just (RenamedPatLet y _) -> do
          -- Use the same variable name
          (stmts', e2) <- toExpE e2
          e <- lift [e|Let $(pure (VarE y)) $(pure t) $(pure e1) $(pure e2)|]
          return (stmts ++ stmts', e)
        Nothing -> do
          -- Introduce a new name
          y <- lift $ fromVarName x
          modify' (\env -> env {vars = (x, RenamedExpLet (AppE var (VarE y))) : vars env})
          (stmts', e2) <- toExpE e2
          e <- lift [e|Let $(pure (VarE y)) $(pure t) $(pure e1) $(pure e2)|]
          return (stmts ++ BindS (VarP y) genVarName : stmts', e)
        _ -> fail $ "Jikka.Core.Language.QuasiRules.toExpE: variable conflicts: " ++ formatVarName x
    Assert e1 e2 -> do
      (stmts1, e1) <- toExpE e1
      (stmts2, e2) <- toExpE e2
      e <- lift [e|Assert $(pure e1) $(pure e2)|]
      return (stmts1 ++ stmts2, e)

alphaExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m Expr
alphaExpr = Alpha.runExpr

ruleExp :: String -> Q Exp
ruleExp s = do
  (name, args, e1, e2) <- liftError $ parseRule s
  (args, e1, e2) <- liftError $ TypeInfer.runRule args e1 e2
  env <-
    return $
      Env
        { vars = map (second (const DeclaredAtForall)) args,
          typeVars = []
        }
  (pat, env) <- runStateT (toPatE e1) env
  supressUnusedMatchesWarnings <- (concat <$>) . forM (vars env) $ \case
    (_, expFromRenamedVarName -> Just e) -> do
      e <- [e|return $(pure e)|]
      return [NoBindS e]
    _ -> return []
  supressUnusedMatchesWarnings' <- forM (typeVars env) $ \(_, y) -> do
    NoBindS <$> [e|return $(pure (VarE y))|]
  ((stmts, exp), _) <- runStateT (toExpE e2) env
  nop <- [e|return ()|]
  exp' <- [e|return $(pure exp)|]
  [e|
    makeRewriteRule $(pure (LitE (StringL name))) $ \env e -> case e of
      $(pure pat) -> do
        $(pure (DoE (supressUnusedMatchesWarnings ++ [NoBindS nop])))
        $(pure (DoE (supressUnusedMatchesWarnings' ++ [NoBindS nop])))
        e <- $(pure (DoE (stmts ++ [NoBindS exp'])))
        Just <$> alphaExpr (typeEnv env) e
      _ -> return Nothing
    |]

r :: TH.QuasiQuoter
r =
  TH.QuasiQuoter
    { TH.quoteExp = ruleExp,
      TH.quotePat = undefined,
      TH.quoteType = undefined,
      TH.quoteDec = undefined
    }