module Jikka.Python.Language.Util where

import Control.Monad.Identity
import Jikka.Common.Location
import Jikka.Python.Language.Expr

constIntExp :: Integer -> Expr
constIntExp = Constant . ConstInt

constBoolExp :: Bool -> Expr
constBoolExp = Constant . ConstBool

mapExprArgumentsM :: Monad m => (Expr' -> m Expr') -> Arguments -> m Arguments
mapExprArgumentsM f args = do
  kwDefaults <- mapM (mapExprM f) (argsKwDefaults args)
  defaults <- mapM (mapExprM f) (argsDefaults args)
  return $
    args
      { argsKwDefaults = kwDefaults,
        argsDefaults = defaults
      }

mapExprComprehensionM :: Monad m => (Expr' -> m Expr') -> Comprehension -> m Comprehension
mapExprComprehensionM f comp = do
  iter <- mapExprM f (compIter comp)
  target <- mapExprM f (compIter comp)
  ifs <- mapM (mapExprM f) (compIfs comp)
  return $ Comprehension {compTarget = target, compIter = iter, compIfs = ifs}

mapExprComprehensionsM :: Monad m => (Expr' -> m Expr') -> [Comprehension] -> m [Comprehension]
mapExprComprehensionsM f = mapM (mapExprComprehensionM f)

mapExprKeywordsM :: Monad m => (Expr' -> m Expr') -> [Keyword'] -> m [Keyword']
mapExprKeywordsM f kwargs = mapM (\(WithLoc loc (k, v)) -> WithLoc loc . (,) k <$> mapExprM f v) kwargs

mapExprM :: Monad m => (Expr' -> m Expr') -> Expr' -> m Expr'
mapExprM f = go
  where
    go e0 =
      f . WithLoc (loc e0) =<< case value e0 of
        BoolOp e1 op e2 -> BoolOp <$> go e1 <*> pure op <*> go e2
        NamedExpr x e -> NamedExpr <$> go x <*> go e
        BinOp e1 op e2 -> BinOp <$> go e1 <*> pure op <*> go e2
        UnaryOp op e -> UnaryOp op <$> go e
        Lambda args body -> Lambda <$> mapExprArgumentsM f args <*> mapExprM f body
        IfExp e1 e2 e3 -> IfExp <$> go e1 <*> go e2 <*> go e3
        Dict es -> Dict <$> mapM (\(k, v) -> (,) <$> traverse go k <*> go v) es
        Set es -> Set <$> mapM go es
        ListComp e comps -> ListComp <$> mapExprM f e <*> mapExprComprehensionsM f comps
        SetComp e comps -> SetComp <$> mapExprM f e <*> mapExprComprehensionsM f comps
        DictComp k v comps -> DictComp <$> mapExprM f k <*> mapExprM f v <*> mapExprComprehensionsM f comps
        GeneratorExp e comps -> GeneratorExp <$> mapExprM f e <*> mapExprComprehensionsM f comps
        Await e -> Await <$> go e
        Yield e -> Yield <$> traverse go e
        YieldFrom e -> YieldFrom <$> go e
        Compare e es -> Compare <$> go e <*> mapM (\(op, e) -> (,) op <$> go e) es
        Call g args kwargs -> Call <$> go g <*> mapM go args <*> mapExprKeywordsM f kwargs
        FormattedValue e1 c e2 -> FormattedValue <$> go e1 <*> pure c <*> traverse go e2
        JoinedStr es -> JoinedStr <$> mapM go es
        Constant constant -> return $ Constant constant
        Attribute e x -> Attribute <$> go e <*> pure x
        Subscript e1 e2 -> Subscript <$> go e1 <*> go e2
        Starred e -> Starred <$> go e
        Name x -> return $ Name x
        List es -> List <$> mapM go es
        Tuple es -> Tuple <$> mapM go es
        Slice e1 e2 e3 -> Slice <$> traverse go e1 <*> traverse go e2 <*> traverse go e3

mapExprExceptHanderM :: Monad m => (Expr' -> m Expr') -> ExceptHandler' -> m ExceptHandler'
mapExprExceptHanderM f (WithLoc loc handler) = do
  body <- mapExprStatementsM f (exchBody handler)
  return $ WithLoc loc (handler {exchBody = body})

mapExprStatementM :: Monad m => (Expr' -> m Expr') -> Statement' -> m Statement'
mapExprStatementM f stmt =
  WithLoc (loc stmt) <$> case value stmt of
    FunctionDef g args body decorators ret -> FunctionDef g args <$> mapExprStatementsM f body <*> mapM (mapExprM f) decorators <*> pure ret
    AsyncFunctionDef g args body decorators ret -> AsyncFunctionDef g args <$> mapExprStatementsM f body <*> mapM (mapExprM f) decorators <*> pure ret
    ClassDef name bases keywords body decorators -> ClassDef name <$> mapM (mapExprM f) bases <*> mapExprKeywordsM f keywords <*> mapExprStatementsM f body <*> mapM (mapExprM f) decorators
    Return e -> Return <$> traverse (mapExprM f) e
    Delete xs -> Delete <$> mapM (mapExprM f) xs
    Assign xs e -> Assign <$> mapM (mapExprM f) xs <*> mapExprM f e
    AugAssign x op e -> AugAssign <$> mapExprM f x <*> pure op <*> mapExprM f e
    AnnAssign x t e -> AnnAssign <$> mapExprM f x <*> mapExprM f t <*> traverse (mapExprM f) e
    For x iter body orelse -> For <$> mapExprM f x <*> mapExprM f iter <*> mapExprStatementsM f body <*> mapExprStatementsM f orelse
    AsyncFor x iter body orelse -> AsyncFor <$> mapExprM f x <*> mapExprM f iter <*> mapExprStatementsM f body <*> mapExprStatementsM f orelse
    While e body orelse -> While <$> mapExprM f e <*> mapExprStatementsM f body <*> mapExprStatementsM f orelse
    If e body orelse -> If <$> mapExprM f e <*> mapExprStatementsM f body <*> mapExprStatementsM f orelse
    With withitems body -> With <$> mapM (\(e1, e2) -> (,) <$> mapExprM f e1 <*> traverse (mapExprM f) e2) withitems <*> mapExprStatementsM f body
    AsyncWith withitems body -> AsyncWith <$> mapM (\(e1, e2) -> (,) <$> mapExprM f e1 <*> traverse (mapExprM f) e2) withitems <*> mapExprStatementsM f body
    Raise e1 e2 -> Raise <$> traverse (mapExprM f) e1 <*> traverse (mapExprM f) e2
    Try body handlers orelse finalbody -> Try <$> mapExprStatementsM f body <*> mapM (mapExprExceptHanderM f) handlers <*> mapExprStatementsM f orelse <*> mapExprStatementsM f finalbody
    Assert e1 e2 -> Assert <$> mapExprM f e1 <*> traverse (mapExprM f) e2
    Import aliases -> return $ Import aliases
    ImportFrom xs aliases -> return $ ImportFrom xs aliases
    Global xs -> return $ Global xs
    Nonlocal xs -> return $ Nonlocal xs
    Expr' e -> Expr' <$> mapExprM f e
    Pass -> return Pass
    Break -> return Break
    Continue -> return Continue

mapExprStatementsM :: Monad m => (Expr' -> m Expr') -> [Statement'] -> m [Statement']
mapExprStatementsM f = mapM (mapExprStatementM f)

mapExprProgramM :: Monad m => (Expr' -> m Expr') -> Program -> m Program
mapExprProgramM = mapExprStatementsM

mapExprProgram :: (Expr' -> Expr') -> Program -> Program
mapExprProgram f = runIdentity . mapExprStatementsM (return . f)