content: {-# LANGUAGE LambdaCase #-}
  
  -- |
  -- Module      : Jikka.RestrictedPython.Format
  -- Description : converts AST of the restricted Python to strings. / 制限された Python の抽象構文木を文字列に変換します。
  -- Copyright   : (c) Kimiyuki Onaka, 2020
  -- License     : Apache License 2.0
  -- Maintainer  : kimiyuki95@gmail.com
  -- Stability   : experimental
  -- Portability : portable
  --
  -- TODO: add parens with considering precedences.
  module Jikka.RestrictedPython.Format
    ( run,
      run',
      formatType,
      formatOperator,
      formatBuiltin,
      formatAttribute,
      formatExpr,
      formatTarget,
    )
  where
  
  import Data.List (intercalate)
  import Data.Text (Text, pack)
  import Jikka.Common.Format.AutoIndent
  import Jikka.RestrictedPython.Language.Builtin
  import Jikka.RestrictedPython.Language.Expr
  
  formatType :: Type -> String
  formatType t = case t of
    VarTy x -> formatTypeName x
    IntTy -> "int"
    BoolTy -> "bool"
    ListTy t -> "List[" ++ formatType t ++ "]"
    NoneTy -> "None"
    TupleTy ts -> "Tuple[" ++ intercalate ", " (map formatType ts) ++ "]"
    CallableTy ts ret -> "Callable[[" ++ intercalate ", " (map formatType ts) ++ "], " ++ formatType ret ++ "]"
    StringTy -> "str"
    SideEffectTy -> "side-effect"
  
  formatConstant :: Constant -> String
  formatConstant = \case
    ConstNone -> "None"
    ConstInt n -> show n
    ConstBool p -> show p
    ConstBuiltin b -> formatBuiltin b
  
  formatBoolOp :: BoolOp -> String
  formatBoolOp = \case
    And -> "and"
    Or -> "or"
    Implies -> "implies"
  
  formatOperator :: Operator -> String
  formatOperator = \case
    Add -> "+"
    Sub -> "-"
    Mult -> "*"
    MatMult -> "@"
    Div -> "/"
    FloorDiv -> "//"
    FloorMod -> "%"
    CeilDiv -> "/^"
    CeilMod -> "%^"
    Pow -> "**"
    BitLShift -> "<<"
    BitRShift -> ">>"
    BitOr -> "|"
    BitXor -> "^"
    BitAnd -> "&"
    Max -> ">?"
    Min -> "<?"
  
  formatUnaryOp :: UnaryOp -> String
  formatUnaryOp = \case
    Invert -> "~"
    Not -> "not"
    UAdd -> "+"
    USub -> "-"
  
  formatCmpOp :: CmpOp' -> String
  formatCmpOp (CmpOp' op _) = case op of
    Eq' -> "=="
    NotEq -> "!="
    Lt -> "<"
    LtE -> "<="
    Gt -> ">"
    GtE -> ">="
    Is -> "is"
    IsNot -> "is not"
    In -> "in"
    NotIn -> "not in"
  
  formatComprehension :: Comprehension -> String
  formatComprehension (Comprehension x iter ifs) =
    let body = "for " ++ formatTarget x ++ " in " ++ formatExpr iter
        ifs' = case ifs of
          Nothing -> ""
          Just ifs -> " if " ++ formatExpr ifs
     in body ++ ifs'
  
  formatTarget :: Target' -> String
  formatTarget (WithLoc' _ x) = case x of
    SubscriptTrg x e -> formatTarget x ++ "[" ++ formatExpr e ++ "]"
    NameTrg x -> formatVarName (value' x)
    TupleTrg xs -> case xs of
      [] -> "()"
      [x] -> "(" ++ formatTarget x ++ ",)"
      _ -> intercalate ", " (map formatTarget xs)
  
  formatExpr :: Expr' -> String
  formatExpr (WithLoc' _ e0) = case e0 of
    BoolOp e1 op e2 -> formatExpr e1 ++ " " ++ formatBoolOp op ++ " " ++ formatExpr e2
    BinOp e1 op e2 -> formatExpr e1 ++ " " ++ formatOperator op ++ " " ++ formatExpr e2
    UnaryOp op e -> formatUnaryOp op ++ " " ++ formatExpr e
    Lambda args body -> case args of
      [] -> "lambda: " ++ formatExpr body
      _ -> "lambda " ++ intercalate ", " (map (formatVarName . value' . fst) args) ++ ": " ++ formatExpr body
    IfExp e1 e2 e3 -> formatExpr e2 ++ " if " ++ formatExpr e1 ++ " else " ++ formatExpr e3
    ListComp e comp -> "[" ++ formatExpr e ++ " " ++ formatComprehension comp ++ "]"
    Compare e1 op e2 -> formatExpr e1 ++ " " ++ formatCmpOp op ++ " " ++ formatExpr e2
    Call f args -> case args of
      [WithLoc' _ (ListComp e comp)] -> formatExpr f ++ "(" ++ formatExpr e ++ " " ++ formatComprehension comp ++ ")"
      _ -> formatExpr f ++ "(" ++ intercalate ", " (map formatExpr args) ++ ")"
    Constant const -> formatConstant const
    Attribute e (WithLoc' _ x) -> formatExpr e ++ "." ++ formatAttribute x
    Subscript e1 e2 -> formatExpr e1 ++ "[" ++ formatExpr e2 ++ "]"
    Starred e -> "*" ++ formatExpr e
    Name x -> formatVarName (value' x)
    List _ es -> "[" ++ intercalate ", " (map formatExpr es) ++ "]"
    Tuple es -> case es of
      [] -> "()"
      [e] -> "(" ++ formatExpr e ++ ",)"
      _ -> "(" ++ intercalate ", " (map formatExpr es) ++ ")"
    SubscriptSlice e from to step ->
      let from' = maybe "" formatExpr from
          to' = maybe "" formatExpr to
          step' = maybe "" ((':' :) . formatExpr) step
       in formatExpr e ++ "[" ++ from' ++ ":" ++ to' ++ step' ++ "]"
  
  formatStatement :: Statement -> [String]
  formatStatement = \case
    Return e -> ["return " ++ formatExpr e]
    AugAssign x op e -> [formatTarget x ++ " " ++ formatOperator op ++ "= " ++ formatExpr e]
    AnnAssign x t e -> [formatTarget x ++ ": " ++ formatType t ++ " = " ++ formatExpr e]
    For x iter body -> ["for " ++ formatTarget x ++ " in " ++ formatExpr iter ++ ":", indent] ++ concatMap formatStatement body ++ [dedent]
    If e body1 body2 -> case body2 of
      [] -> ["if " ++ formatExpr e ++ ":", indent] ++ concatMap formatStatement body1 ++ [dedent]
      [body2@(If _ _ _)] ->
        let elif : cont = formatStatement body2
         in ["if " ++ formatExpr e ++ ":", indent] ++ concatMap formatStatement body1 ++ [dedent, "el" ++ elif] ++ cont
      _ -> ["if " ++ formatExpr e ++ ":", indent] ++ concatMap formatStatement body1 ++ [dedent, "else:", indent] ++ concatMap formatStatement body2 ++ [dedent]
    Assert e -> ["assert " ++ formatExpr e]
    Expr' e -> [formatExpr e]
  
  formatToplevelStatement :: ToplevelStatement -> [String]
  formatToplevelStatement = \case
    ToplevelAnnAssign x t e -> [formatVarName (value' x) ++ ": " ++ formatType t ++ " = " ++ formatExpr e]
    ToplevelFunctionDef f args ret body -> ["def " ++ formatVarName (value' f) ++ "(" ++ intercalate ", " (map (\(x, t) -> formatVarName (value' x) ++ ": " ++ formatType t) args) ++ ") -> " ++ formatType ret ++ ":", indent] ++ concatMap formatStatement body ++ [dedent]
    ToplevelAssert e -> ["assert " ++ formatExpr e]
  
  formatProgram :: Program -> [String]
  formatProgram prog = concatMap formatToplevelStatement prog
  
  run' :: Program -> String
  run' = unlines . makeIndentFromMarkers 4 . formatProgram
  
  run :: Applicative m => Program -> m Text
  run = pure . pack . run'
  