{-# 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.Monad.Reader
import Data.Maybe
import qualified Data.IntMap as IntMap

import Agda.Syntax.Common
import Agda.Syntax.Common.Pretty
import Agda.Syntax.Treeless

import Agda.Compiler.Treeless.Subst

import Agda.Utils.Function
import Agda.Utils.List
import Agda.Utils.Null (empty)
import Agda.Utils.Tuple (first)

import Agda.Utils.Impossible

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)

instance Pretty TAlt where
  prettyPrec p b = runP $ prec p (pAlt b)

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 [(strictness s <> text x) <+> "=", nest 2 e]
                 | (s, x, e) <- binds
                 ] <+>
            "in"
          , b
          ])
      <$> pLets (zip xs es) b
    where
      (es, b) = tLetView t

      pLets [] b = ([],) <$> pTerm' 0 b
      pLets ((x, (s, e)) : bs) b = do
        e <- pTerm' 0 e
        first ((s, x, e) :) <$> bindName x (pLets bs b)

      strictness Strict    = "!"
      strictness NonStrict = "~"

  TCase x i def alts -> paren 0 $
    (\ sc alts defd ->
      sep [ "case" <+>
            (if isErased (caseErased i) then "(erased)" else empty) <+>
            sc <+> "of"
          , nest 2 $ vcat (alts ++ [ "_ →" <+> defd | null alts || def /= TError TUnreachable ]) ]
    ) <$> pTerm' 0 (TVar x)
      <*> mapM pAlt alts
      <*> pTerm' 0 def

  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

pAlt' :: Doc -> Doc -> Doc
pAlt' p b = sep [p <+> "→", nest 2 b]

pAlt :: TAlt -> P Doc
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