content: {-# OPTIONS_GHC -Wunused-imports #-}
  
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  module Agda.Compiler.Treeless.Pretty () where
  
  import Prelude hiding ((!!)) -- don't use partial functions!
  
  import Control.Arrow (first)
  import Control.Monad.Reader
  import Data.Maybe
  import qualified Data.IntMap as IntMap
  
  import Agda.Syntax.Treeless
  import Agda.Syntax.Common.Pretty
  
  import Agda.Compiler.Treeless.Subst
  
  import Agda.Utils.Impossible
  import Agda.Utils.Function
  import Agda.Utils.List
  
  instance Pretty Compiled where
    pretty Compiled {cTreeless, cArgUsage} =
      "Compiled {" <?> vcat
        [ "cTreeless   =" <?> pretty cTreeless
        , "funCompiled =" <?> pshow cArgUsage
        ] <?> "}"
  
  data PEnv = PEnv { pPrec :: Int
                   , pFresh :: [String]
                   , pBound :: [String] }
  
  type P = Reader PEnv
  
  --UNUSED Liang-Ting Chen 2019-07-16
  --withName :: (String -> P a) -> P a
  --withName k = withNames 1 $ \[x] -> k x
  
  withNames :: Int -> ([String] -> P a) -> P a
  withNames n k = do
    (xs, ys) <- asks $ splitAt n . pFresh
    local (\ e -> e { pFresh = ys }) (k xs)
  
  -- | Don't generate fresh names for unused variables.
  withNames' :: HasFree a => Int -> a -> ([String] -> P b) -> P b
  withNames' n tm k = withNames n' $ k . insBlanks
    where
      fv = freeVars tm
      n'  = length $ filter (< n) $ IntMap.keys fv
      insBlanks = go n
        where
          go 0 _ = []
          go i xs0@(~(x : xs))
            | IntMap.member (i - 1) fv = x   : go (i - 1) xs
            | otherwise             = "_" : go (i - 1) xs0
  
  bindName :: String -> P a -> P a
  bindName x = local $ \ e -> e { pBound = x : pBound e }
  
  bindNames :: [String] -> P a -> P a
  bindNames xs p = foldr bindName p xs
  
  paren :: Int -> P Doc -> P Doc
  paren p doc = do
    n <- asks pPrec
    applyWhen (p < n) parens <$> doc
  
  prec :: Int -> P a -> P a
  prec p = local $ \ e -> e { pPrec = p }
  
  name :: Int -> P String
  name x = asks
    $ (\ xs -> indexWithDefault __IMPOSSIBLE__ xs x)
    . (++ map (("^" ++) . show) [1..])
    . pBound
  
  runP :: P a -> a
  runP p = runReader p PEnv{ pPrec = 0, pFresh = names, pBound = [] }
    where
      names = [ x ++ i | i <- "" : map show [1..], x <- map (:[]) ['a'..'z'] ]
  
  instance Pretty TTerm where
    prettyPrec p t = runP $ prec p (pTerm t)
  
  opName :: TPrim -> String
  opName PAdd = "+"
  opName PSub = "-"
  opName PMul = "*"
  opName PQuot = "quot"
  opName PRem = "rem"
  opName PGeq = ">="
  opName PLt  = "<"
  opName PEqI = "==I"
  opName PAdd64 = "+64"
  opName PSub64 = "-64"
  opName PMul64 = "*64"
  opName PQuot64 = "quot64"
  opName PRem64 = "rem64"
  opName PLt64  = "<64"
  opName PEq64 = "==64"
  opName PEqF = "==F"
  opName PEqS = "==S"
  opName PEqC = "==C"
  opName PEqQ = "==Q"
  opName PIf  = "if_then_else_"
  opName PSeq = "seq"
  opName PITo64 = "toWord64"
  opName P64ToI = "fromWord64"
  
  
  isInfix :: TPrim -> Maybe (Int, Int, Int)
  isInfix op =
    case op of
      PMul -> l 7
      PAdd -> l 6
      PSub -> l 6
      PGeq -> non 4
      PLt  -> non 4
      PMul64 -> l 7
      PAdd64 -> l 6
      PSub64 -> l 6
      PLt64  -> non 4
      p | isPrimEq p -> non 4
      _    -> Nothing
    where
      l n   = Just (n, n, n + 1)
      r n   = Just (n, n + 1, n) -- NB:: Defined but not used
      non n = Just (n, n + 1, n + 1)
  
  pTerm' :: Int -> TTerm -> P Doc
  pTerm' p = prec p . pTerm
  
  pTerm :: TTerm -> P Doc
  pTerm = \case
    TVar x -> text <$> name x
    TApp (TPrim op) [a, b] | Just (c, l, r) <- isInfix op ->
      paren c $ sep <$> sequence [ pTerm' l a
                                 , pure $ text $ opName op
                                 , pTerm' r b ]
    TApp (TPrim PIf) [a, b, c] ->
      paren 0 $ (\ a b c -> sep [ "if" <+> a
                                , nest 2 $ "then" <+> b
                                , nest 2 $ "else" <+> c ])
                <$> pTerm' 0 a
                <*> pTerm' 0 b
                <*> pTerm c
    TDef f -> pure $ pretty f
    TCon c -> pure $ pretty c
    TLit l -> pure $ pretty l
    TPrim op | isJust (isInfix op) -> pure $ text ("_" ++ opName op ++ "_")
             | otherwise -> pure $ text (opName op)
    TApp f es ->
      paren 9 $ (\a bs -> sep [a, nest 2 $ fsep bs])
                <$> pTerm' 9 f
                <*> mapM (pTerm' 10) es
    t@TLam{} -> paren 0 $ withNames' n b $ \ xs -> bindNames xs $
      (\b -> sep [ text ("λ " ++ unwords xs ++ " →")
                 , nest 2 b ]) <$> pTerm' 0 b
      where
        (n, b) = tLamView t
    t@TLet{} -> paren 0 $ withNames (length es) $ \ xs ->
      (\ (binds, b) -> sep [ "let" <+> vcat [ sep [ text x <+> "="
                                                  , nest 2 e ] | (x, e) <- binds ]
                                <+> "in", b ])
        <$> pLets (zip xs es) b
      where
        (es, b) = tLetView t
  
        pLets [] b = ([],) <$> pTerm' 0 b
        pLets ((x, e) : bs) b = do
          e <- pTerm' 0 e
          first ((x, e) :) <$> bindName x (pLets bs b)
  
    TCase x _ def alts -> paren 0 $
      (\ sc alts defd ->
        sep [ "case" <+> sc <+> "of"
            , nest 2 $ vcat (alts ++ [ "_ →" <+> defd | null alts || def /= TError TUnreachable ]) ]
      ) <$> pTerm' 0 (TVar x)
        <*> mapM pAlt alts
        <*> pTerm' 0 def
      where
        pAlt (TALit l b) = pAlt' <$> pTerm' 0 (TLit l) <*> pTerm' 0 b
        pAlt (TAGuard g b) =
          pAlt' <$> (("_" <+> "|" <+>) <$> pTerm' 0 g)
                <*> (pTerm' 0 b)
        pAlt (TACon c a b) =
          withNames' a b $ \ xs -> bindNames xs $
          pAlt' <$> pTerm' 0 (TApp (TCon c) [TVar i | i <- reverse [0..a - 1]])
                <*> pTerm' 0 b
        pAlt' p b = sep [p <+> "→", nest 2 b]
  
    TUnit -> pure "()"
    TSort -> pure "Set"
    TErased -> pure "_"
    TError err -> paren 9 $ pure $ "error" <+> text (show (show err))
    TCoerce t -> paren 9 $ ("coe" <+>) <$> pTerm' 10 t
  