{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}

-- |
-- Module      : Jikka.Core.Format
-- Description : converts the syntax trees of core language to strings. / core 言語の構文木を文字列に変換します。
-- 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.Core.Format
  ( run,
    formatBuiltinIsolated,
    formatBuiltin,
    formatType,
    formatExpr,
    formatToplevelExpr,
    formatProgram,
  )
where

import Data.Char (toLower)
import Data.List (intercalate)
import Data.Text (Text, pack)
import Jikka.Common.Format.AutoIndent
import Jikka.Core.Language.BuiltinPatterns (pattern Range1')
import Jikka.Core.Language.Expr
import Jikka.Core.Language.FreeVars (isUnusedVar)
import Jikka.Core.Language.LambdaPatterns
import Jikka.Core.Language.Util

-- | See also Table 2 of <https://www.haskell.org/onlinereport/decls.html Haskell Online Report, 4 Declarations and Bindings>.
newtype Prec = Prec Int
  deriving (Eq, Ord, Show, Read)

instance Enum Prec where
  toEnum n = Prec n
  fromEnum (Prec n) = n

identPrec = Prec 12

funCallPrec = Prec 11

unaryPrec = Prec 10

powerPrec = Prec 8

multPrec = Prec 7

addPrec = Prec 6

appendPrec = Prec 5

comparePrec = Prec 4

andPrec = Prec 3

orPrec = Prec 2

impliesPrec = Prec 1

commaPrec = Prec 0

lambdaPrec = Prec (-1)

parenPrec = Prec (-2)

data Assoc
  = NoAssoc
  | LeftToRight
  | RightToLeft
  deriving (Eq, Ord, Enum, Show, Read)

paren :: String -> String
paren s = "(" ++ s ++ ")"

-- | `resolvePrec` inserts parens to the given string if required.
--
-- >>> resolvePrec multPrec ("1 + 2", addPrec) ++ " * 3"
-- "(1 + 2) * 3"
--
-- >>> resolvePrec addPrec ("1 * 2", multPrec) ++ " + 3"
-- "1 * 2 + 3"
resolvePrec :: Prec -> (String, Prec) -> String
resolvePrec cur (s, prv)
  | cur > prv = paren s
  | otherwise = s

-- | `resolvePrecLeft` inserts parens to the given string if required.
--
-- >>> resolvePrecLeft addPrec LeftToRight ("1 - 2", addPrec) ++ " - 3"
-- "1 - 2 - 3"
resolvePrecLeft :: Prec -> Assoc -> (String, Prec) -> String
resolvePrecLeft cur assoc (s, prv)
  | cur > prv || (cur == prv && assoc /= LeftToRight) = paren s
  | otherwise = s

-- | `resolvePrecRight` inserts parens to the given string if required.
--
-- >>> "1 - " ++ resolvePrecRight addPrec LeftToRight ("2 - 3", addPrec)
-- "1 - (2 - 3)"
resolvePrecRight :: Prec -> Assoc -> (String, Prec) -> String
resolvePrecRight cur assoc (s, prv)
  | cur > prv || (cur == prv && assoc /= RightToLeft) = paren s
  | otherwise = s

formatType' :: Type -> (String, Prec)
formatType' = \case
  VarTy a -> (formatTypeName a, identPrec)
  IntTy -> ("int", identPrec)
  BoolTy -> ("bool", identPrec)
  ListTy t -> (resolvePrec funCallPrec (formatType' t) ++ " list", funCallPrec)
  TupleTy ts -> case ts of
    [] -> ("unit", identPrec)
    [t] -> (resolvePrec funCallPrec (formatType' t) ++ " one_tuple", funCallPrec)
    _ -> (intercalate " * " (map (resolvePrec (pred multPrec) . formatType') ts), multPrec)
  FunTy t1 t2 ->
    (resolvePrecLeft impliesPrec RightToLeft (formatType' t1) ++ " -> " ++ resolvePrecRight impliesPrec RightToLeft (formatType' t2), impliesPrec)
  DataStructureTy ds -> (formatDataStructure ds, identPrec)

formatType :: Type -> String
formatType = resolvePrec parenPrec . formatType'

formatDataStructure :: DataStructure -> String
formatDataStructure = \case
  ConvexHullTrick -> "convex_hull_trick"
  SegmentTree semigrp -> "segment_tree<" ++ formatSemigroup semigrp ++ ">"

formatSemigroup :: Semigroup' -> String
formatSemigroup = \case
  SemigroupIntPlus -> "int_plus"
  SemigroupIntMin -> "int_min"
  SemigroupIntMax -> "int_max"
  SemigroupIntGcd -> "int_gcd"
  SemigroupIntLcm -> "int_lcm"

data Builtin'
  = Fun String
  | PrefixOp String
  | InfixOp String Prec Assoc
  | At'
  | SetAt'
  | Tuple'
  | Proj' Integer
  | If'
  deriving (Eq, Ord, Show, Read)

funMat :: String -> [Integer] -> Builtin'
funMat f args = Fun $ intercalate "@" (f : map show args)

analyzeBuiltin :: Builtin -> Builtin'
analyzeBuiltin = \case
  -- arithmetical functions
  Negate -> PrefixOp "-"
  Plus -> InfixOp "+" addPrec LeftToRight
  Minus -> InfixOp "-" addPrec LeftToRight
  Mult -> InfixOp "*" multPrec LeftToRight
  FloorDiv -> InfixOp "/" multPrec LeftToRight
  FloorMod -> InfixOp "%" multPrec LeftToRight
  CeilDiv -> InfixOp "/^" multPrec LeftToRight
  CeilMod -> InfixOp "%^" multPrec LeftToRight
  JustDiv -> InfixOp "/!" multPrec LeftToRight
  Pow -> InfixOp "**" powerPrec RightToLeft
  -- advanced arithmetical functions
  Abs -> Fun "abs"
  Gcd -> Fun "gcd"
  Lcm -> Fun "lcm"
  Min2 -> Fun "min"
  Max2 -> Fun "max"
  -- logical functions
  Not -> Fun "not"
  And -> InfixOp "&&" andPrec RightToLeft
  Or -> InfixOp "||" orPrec RightToLeft
  Implies -> Fun "implies"
  If -> If'
  -- bitwise functions
  BitNot -> PrefixOp "~"
  BitAnd -> InfixOp "&" multPrec LeftToRight
  BitOr -> InfixOp "|" appendPrec LeftToRight
  BitXor -> InfixOp "^" addPrec LeftToRight
  BitLeftShift -> InfixOp "<<" powerPrec LeftToRight
  BitRightShift -> InfixOp ">>" powerPrec LeftToRight
  -- matrix functions
  MatAp h w -> funMat "matap" [h, w]
  MatZero h w -> funMat "matzero" [h, w]
  MatOne n -> funMat "matone" [n]
  MatAdd h w -> funMat "matadd" [h, w]
  MatMul h n w -> funMat "matmul" [h, n, w]
  MatPow n -> funMat "matpow" [n]
  VecFloorMod n -> funMat "vecfloormod" [n]
  MatFloorMod h w -> funMat "matfloormod" [h, w]
  -- modular functions
  ModNegate -> Fun "modnegate"
  ModPlus -> Fun "modplus"
  ModMinus -> Fun "modminus"
  ModMult -> Fun "modmult"
  ModInv -> Fun "modinv"
  ModPow -> Fun "modpow"
  ModMatAp h w -> funMat "modmatap" [h, w]
  ModMatAdd h w -> funMat "modmatadd" [h, w]
  ModMatMul h n w -> funMat "modmatmul" [h, n, w]
  ModMatPow n -> funMat "modmatpow" [n]
  -- list functions
  Cons -> Fun "cons"
  Snoc -> Fun "snoc"
  Foldl -> Fun "foldl"
  Scanl -> Fun "scanl"
  Build -> Fun "build"
  Iterate -> Fun "iterate"
  Len -> Fun "len"
  Map -> Fun "map"
  Filter -> Fun "filter"
  At -> At'
  SetAt -> SetAt'
  Elem -> Fun "elem"
  Sum -> Fun "sum"
  Product -> Fun "product"
  ModSum -> Fun "modsum"
  ModProduct -> Fun "modproduct"
  Min1 -> Fun "minimum"
  Max1 -> Fun "maximum"
  ArgMin -> Fun "argmin"
  ArgMax -> Fun "argmax"
  Gcd1 -> Fun "gcds"
  Lcm1 -> Fun "lcms"
  All -> Fun "all"
  Any -> Fun "any"
  Sorted -> Fun "sort"
  Reversed -> Fun "reverse"
  Range1 -> Fun "range"
  Range2 -> Fun "range2"
  Range3 -> Fun "range3"
  -- tuple functions
  Tuple -> Tuple'
  Proj n -> Proj' n
  -- comparison
  LessThan -> InfixOp "<" comparePrec NoAssoc
  LessEqual -> InfixOp "<=" comparePrec NoAssoc
  GreaterThan -> InfixOp ">" comparePrec NoAssoc
  GreaterEqual -> InfixOp ">=" comparePrec NoAssoc
  Equal -> InfixOp "==" comparePrec NoAssoc
  NotEqual -> InfixOp "!=" comparePrec NoAssoc
  -- combinational functions
  Fact -> Fun "fact"
  Choose -> Fun "choose"
  Permute -> Fun "permute"
  MultiChoose -> Fun "multichoose"
  -- data structures
  ConvexHullTrickInit -> Fun "cht_init"
  ConvexHullTrickGetMin -> Fun "cht_getmin"
  ConvexHullTrickInsert -> Fun "cht_insert"
  SegmentTreeInitList _ -> Fun "segtree_initlist"
  SegmentTreeGetRange _ -> Fun "segtree_getrange"
  SegmentTreeSetPoint _ -> Fun "segtree_setpoint"

formatTemplate :: [Type] -> String
formatTemplate ts = concatMap (('@' :) . formatType) ts

formatFunCall :: (String, Prec) -> [Expr] -> (String, Prec)
formatFunCall f [] = f
formatFunCall f args =
  let f' = resolvePrecLeft funCallPrec LeftToRight f
      args' = map (resolvePrecRight funCallPrec LeftToRight . formatExpr') args
   in (unwords (f' : args'), funCallPrec)

formatBuiltinIsolated' :: Builtin' -> [Type] -> String
formatBuiltinIsolated' builtin ts = case builtin of
  Fun name -> name ++ formatTemplate ts
  PrefixOp op -> paren $ op ++ formatTemplate ts
  InfixOp op _ _ -> paren $ op ++ formatTemplate ts
  At' -> paren $ "at" ++ formatTemplate ts
  SetAt' -> paren $ "set-at" ++ formatTemplate ts
  Tuple' -> paren $ "tuple" ++ formatTemplate ts
  Proj' n -> paren $ "proj-" ++ show n ++ formatTemplate ts
  If' -> paren $ "if-then-else" ++ formatTemplate ts

formatBuiltinIsolated :: Builtin -> [Type] -> String
formatBuiltinIsolated builtin ts = formatBuiltinIsolated' (analyzeBuiltin builtin) ts

formatBuiltin' :: Builtin -> [Type] -> [Expr] -> (String, Prec)
formatBuiltin' builtin ts args = case (analyzeBuiltin builtin, ts, args) of
  (Fun "map", _, [Lam x IntTy e, Range1' n]) | x `isUnusedVar` e -> formatFunCall ("replicate", identPrec) [n, e]
  (Fun name, _, _) -> formatFunCall (name, identPrec) args
  (PrefixOp op, _, e1 : args) -> formatFunCall (op ++ " " ++ resolvePrec unaryPrec (formatExpr' e1), unaryPrec) args
  (InfixOp op prec assoc, _, e1 : e2 : args) -> formatFunCall (resolvePrecLeft prec assoc (formatExpr' e1) ++ " " ++ op ++ " " ++ resolvePrecRight prec assoc (formatExpr' e2), prec) args
  (At', _, e1 : e2 : args) -> formatFunCall (resolvePrec identPrec (formatExpr' e1) ++ "[" ++ resolvePrec parenPrec (formatExpr' e2) ++ "]", identPrec) args
  (SetAt', _, e1 : e2 : e3 : args) -> formatFunCall (resolvePrec identPrec (formatExpr' e1) ++ "[" ++ resolvePrec parenPrec (formatExpr' e2) ++ " <- " ++ resolvePrec parenPrec (formatExpr' e3) ++ "]", identPrec) args
  (Tuple', [_], e : args) -> formatFunCall (paren (resolvePrec commaPrec (formatExpr' e) ++ ","), identPrec) args
  (Tuple', _, args) | length args >= length ts -> formatFunCall (paren (intercalate ", " (map (resolvePrec commaPrec . formatExpr') (take (length ts) args))), identPrec) (drop (length ts) args)
  (Proj' n, _, e : args) -> formatFunCall (resolvePrec identPrec (formatExpr' e) ++ "." ++ show n, identPrec) args
  (If', _, e1 : e2 : e3 : args) -> formatFunCall ("if" ++ " " ++ resolvePrec parenPrec (formatExpr' e1) ++ " then " ++ resolvePrec parenPrec (formatExpr' e2) ++ " else " ++ resolvePrec lambdaPrec (formatExpr' e3), lambdaPrec) args
  _ -> formatFunCall (formatBuiltinIsolated' (analyzeBuiltin builtin) ts, identPrec) args

formatBuiltin :: Builtin -> [Type] -> [Expr] -> String
formatBuiltin f ts args = resolvePrec parenPrec (formatBuiltin' f ts args)

formatLiteral :: Literal -> String
formatLiteral = \case
  LitBuiltin builtin ts -> formatBuiltinIsolated builtin ts
  LitInt n -> show n
  LitBool p -> map toLower $ show p
  LitNil _ -> "nil"
  LitBottom _ msg -> "bottom<" ++ show msg ++ ">"

formatFormalArgs :: [(VarName, Type)] -> String
formatFormalArgs args = unwords $ map (\(x, t) -> paren (formatVarName x ++ ": " ++ formatType t)) args

formatExpr' :: Expr -> (String, Prec)
formatExpr' = \case
  Var x -> (formatVarName x, identPrec)
  Lit lit -> (formatLiteral lit, identPrec)
  e@(App _ _) ->
    let (f, args) = curryApp e
     in case f of
          Var x -> formatFunCall (formatVarName x, identPrec) args
          Lit (LitBuiltin builtin ts) -> formatBuiltin' builtin ts args
          _ -> formatFunCall (formatExpr' f) args
  LamId _ -> ("id", identPrec)
  LamConst _ e -> formatFunCall ("const", identPrec) [e]
  e@(Lam _ _ _) ->
    let (args, body) = uncurryLam e
     in ("fun " ++ formatFormalArgs args ++ " ->\n" ++ indent ++ "\n" ++ resolvePrec parenPrec (formatExpr' body) ++ "\n" ++ dedent ++ "\n", lambdaPrec)
  Let x t e1 e2 -> ("let " ++ formatVarName x ++ ": " ++ formatType t ++ " = " ++ resolvePrec parenPrec (formatExpr' e1) ++ "\nin " ++ resolvePrec lambdaPrec (formatExpr' e2), lambdaPrec)
  Assert e1 e2 -> ("assert " ++ resolvePrec parenPrec (formatExpr' e1) ++ " in\n" ++ resolvePrec lambdaPrec (formatExpr' e2), lambdaPrec)

formatExpr :: Expr -> String
formatExpr = unlines . makeIndentFromMarkers 4 . lines . resolvePrec parenPrec . formatExpr'

formatToplevelExpr' :: ToplevelExpr -> [String]
formatToplevelExpr' = \case
  ResultExpr e -> lines (resolvePrec lambdaPrec (formatExpr' e))
  ToplevelLet x t e cont -> let' (formatVarName x) t e cont
  ToplevelLetRec f args ret e cont -> let' ("rec " ++ formatVarName f ++ " " ++ formatFormalArgs args) ret e cont
  ToplevelAssert e cont -> ["assert " ++ resolvePrec parenPrec (formatExpr' e), "in"] ++ formatToplevelExpr' cont
  where
    let' s t e cont =
      ["let " ++ s ++ ": " ++ formatType t ++ " =", indent]
        ++ lines (resolvePrec parenPrec (formatExpr' e))
        ++ [dedent, "in"]
        ++ formatToplevelExpr' cont

formatToplevelExpr :: ToplevelExpr -> String
formatToplevelExpr = unlines . makeIndentFromMarkers 4 . formatToplevelExpr'

formatProgram :: Program -> String
formatProgram = formatToplevelExpr

run :: Applicative m => Program -> m Text
run = pure . pack . formatProgram