{-# LANGUAGE TypeSynonymInstances
           , FlexibleInstances
           , ScopedTypeVariables
           , RankNTypes
           #-}
module Spire.Surface.PrettyPrinter (prettyPrint , prettyPrintError) where
import Spire.Canonical.Embedder
import Spire.Canonical.Types
import Spire.Expression.Embedder
import Spire.Expression.Types
import Spire.Surface.Types

import Control.Applicative ((<$>))
import Control.Monad.Reader
import qualified Text.PrettyPrint.Leijen as PP
import Text.PrettyPrint.Leijen (Doc)
import Unbound.LocallyNameless hiding ( Spine )
import Unbound.LocallyNameless.Fresh

----------------------------------------------------------------------

prettyPrint :: Display t => t -> String
prettyPrint t = render $ runFreshM $ runReaderT (display t) emptyDisplayR
  where
  -- To adjust the printing parameters:
  render x = PP.displayS (PP.renderPretty ribbon columns x) ""
  columns = 80
  ribbon = 1.0

-- The raw data is printed in color to make it easier to ignore :P
prettyPrintError :: (Show t , Display t) => t -> String
prettyPrintError a  = prettyPrint a ++ "\x1b[35m(" ++ show a ++ ")\x1b[0m"

----------------------------------------------------------------------

-- Lift standard pretty printing ops to a monad.
sepM , fsepM , hsepM , vcatM , listM , alignSepM , groupNestSepM ::
  (Functor m , Monad m) => [m Doc] -> m Doc
sepM  xs = PP.sep     <$> sequence xs
fsepM xs = PP.fillSep <$> sequence xs
hsepM xs = PP.hsep    <$> sequence xs
vsepM xs = PP.vsep    <$> sequence xs
vcatM xs = PP.vcat    <$> sequence xs
listM xs = PP.list    <$> sequence xs
alignSepM = alignM . sepM
groupNestSepM = groupM . nestM 2 . sepM

nestM , indentM :: Functor m => Int -> m Doc -> m Doc
nestM   n = fmap $ PP.nest n
indentM n = fmap $ PP.indent n

parensM , groupM , alignM :: Functor m => m Doc -> m Doc
parensM = fmap PP.parens
groupM  = fmap PP.group
alignM  = fmap PP.align

infixr 5 </> , <$$> , <$+$>
infixr 6 <> , <+>
(<>) , (<+>) , (</>) , (<$$>) , (<$+$>) ::
  Monad m => m Doc -> m Doc -> m Doc
(<>) = liftM2 (PP.<>)
(<+>) = liftM2 (PP.<+>)
(</>) = liftM2 (PP.</>)
(<$$>) = liftM2 (PP.<$$>)
(<$+$>) = liftM2 (PP.<$>)

dt :: String -> DisplayM Doc
dt = return . PP.text

----------------------------------------------------------------------

binding :: (Precedence t', Precedence t, Display t) =>
           t' -> Nom -> t -> DisplayM Doc
binding outside nm tp | isWildcard nm = wrapNonAssoc outside tp
binding _       nm tp = parensM . hsepM $ [pushName nm (d nm) , dt ":" , d tp]

----------------------------------------------------------------------
-- Short hands.

d :: Display t => t -> DisplayM Doc
d = display

w , ww :: (Precedence o , Precedence i, Display i) =>
          o -> i -> DisplayM Doc
w  = wrap         -- Mnemonic: (w)rapped display
ww = wrapNonAssoc -- Mnemonic: (w)rapped (w)rapped display
                  -- (i.e. more wrapped :P)

----------------------------------------------------------------------
-- Syntactic-category agnostic printers.
--
-- The idea was to write all the pretty printers using these
-- ... except now all but one printer is defined by embedding.

-- Constructors with no arguments have printers with no arguments.
dNil      = dt "[]"
dRefl     = dt "refl"
dHere     = dt "here"
dWildCard = dt wildcard

dQuotes :: String -> DisplayM Doc
dQuotes = return . PP.dquotes . PP.text

-- Constructors with arguments pass *themselves* and their args to
-- their printers.

dCons o x y =  ww o x <+> dt "::" <+> w o y
dPair o x y =  ww o x <+> dt ","  <+> w o y
dEq   o x y =  ww o x <+> dt "==" <+> w o y

dLam o bnd_b = do
  (nm , b) <- unbind bnd_b
  fsepM [ dt "\\" <+> pushName nm (d nm) <+> dt "->" , pushName nm (w o b) ]
dLamArg o bnd_b = parensM (dLam o bnd_b)

dPi o _A bnd_B = do
  (nm , _B) <- unbind bnd_B
  fsepM [ binding o nm _A <+> dt "->" , pushName nm (w o _B) ]
dSg o _A bnd_B = do
  (nm , _B) <- unbind bnd_B
  binding o nm _A <+> dt "*" <+> pushName nm (d _B)

dIf o c t f = groupNestSepM
  [ dt "if" <+> w o c
  , dt "then" <+> w o t
  , dt "else" <+> w o f
  ]

dThere o t   = dt "there" <+> ww o t
dEnd   o i   = dt "End"   <+> ww o i
dTag   o _E  = dt "Tag"   <+> ww o _E
dInit  o xs  = dt "init"  <+> ww o xs

dApp o f a = alignM $ w o f </> ww o a
dAnn _ a _A = parensM . alignSepM $ [ d a , dt ":" <+> d _A ]

dRec o i _D =
  dt "Rec" <+> groupNestSepM
    [ ww o i , ww o _D ]

dArg o _A _B =
  dt "Arg" <+> groupNestSepM
    [ ww o _A , dLamArg o _B ]

----------------------------------------------------------------------

instance Display Syntax where
  display s = case s of
    SRefl  -> dRefl
    SHere  -> dHere

    SThere t   -> dThere s t
    SEnd   i   -> dEnd   s i
    SPair a b  -> dPair  s a b
    SLam  b    -> dLam s b

    SRec      i  _D    -> dRec      s i  _D
    SInit     xs       -> dInit     s xs
    SPi       _A _B    -> dPi       s _A _B
    SSg       _A _B    -> dSg       s _A _B
    SArg      _A _B    -> dArg      s _A _B

    SEq a b   -> dEq s a b

    SVar nm   -> d nm
    SQuotes s -> dQuotes s
    SWildCard -> dWildCard
    SIf c t f -> dIf s c t f
    SApp f a  -> dApp s f a
    SAnn a _A -> dAnn s a _A

instance Display Stmt where
  display (SDef nm a _A) =
    (groupM . nestM 2) (d nm <+> dt ":" <$+$> d _A) <$$>
    (groupM . nestM 2) (d nm <+> dt "=" <$+$> d a)

  display (SData _N _P _I css) = let
      params = map (\(l , _A) -> parensM $ dt l <+> dt ":" <+> d _A ) _P
      indices = PP.hcat . PP.punctuate (PP.text " => ") <$> mapM maybeParens (_I ++ [(wildcard , sType)])
      constrs = vsepM $ map (\ (nm , cs) -> dt nm <+> dt ":" <+> displayConstrs _N (map fst _P) cs) css
    in
    dt "data" <+> hsepM (dt _N : params) <+> dt ":" <+> indices <+> dt "where"
    <$$> indentM 2 constrs <$$>
    dt "end"
   where
   maybeParens = \(l , _A) -> if l == wildcard then d _A else parensM (dt l <+> dt ":" <+> d _A)

   displayConstrs :: String -> [String] -> [Constr] -> DisplayM Doc
   displayConstrs _N _P cs = PP.hcat . PP.punctuate (PP.text " => ") <$> mapM (displayConstr _N _P) cs

   displayConstr :: String -> [String] -> Constr -> DisplayM Doc
   displayConstr _N _P (Fix _I) = hsepM $ (dt _N : map dt _P) ++ map d _I
   displayConstr _N _P (Arg nm _A) = maybeParens (nm , _A)

instance Display SProg where
  display defs =  PP.vcat . PP.punctuate PP.line <$> mapM d defs

instance Display Nom where
  -- It's incorrect to use 'name2String' here, since it allows
  -- shadowing, but using 'show' is really ugly, since it prints many
  -- unnecessary freshenings.
  --
  -- The approach we take is to only freshen a name 'x' if there is
  -- another name 'y' in scope for the binding sight of 'x' s.t. 'x'
  -- and 'y' have the same string part (a 'Name' is essentially a
  -- '(String , Int)' pair, with the int used for freshening).
  --
  -- For example, the function
  --
  --   \ x . \ x . x
  --
  -- will print as
  --
  --   \ x . \ x$<n> . x$<n> .
  --
  -- A much fancier printer could detect that this example need not be
  -- freshened, but we're taking a simple approach here.
  display nm = do
    nms <- asks names
        -- Here 'nms'' is the names in scope for the binding of 'nm'.
        -- The 'nms' is a stack of bindings, so all the bindings after
        -- 'nm' correspond to bindings in scope for 'nm's binding.
        --
        -- Metavars are always freshened.
    let nms'   = drop 1 . dropWhile (/= nm) $ nms
        suffix = if name2String nm `elem` map name2String nms'
                    || isMV nm
                 then "$" ++ show (name2Integer nm)
                 else ""
    dt $ name2String nm ++ suffix

----------------------------------------------------------------------

instance Display CProg where
  display = d . runFreshM . mapM embedCDef

instance Display VProg where
  display = d . runFreshM . mapM (embedCDef <=< embedVDef)

instance Display Check where
  display = d . runFreshM . embedC

instance Display Infer where
  display = d . runFreshM . embedI

----------------------------------------------------------------------

instance Display Value where
  display = d . runFreshM . (embedC <=< embedV)

instance Display (Nom , Spine) where
  display (nm , fs) = d . runFreshM $ (embedI =<< embedN nm fs)

instance Display Tel where
  display = listM . map (uncurry $ dAnn undefined) . tel2List

----------------------------------------------------------------------

instance Precedence Syntax where
  level s = case s of
    SPair _ _           -> pairLevel
    SEq _ _             -> eqLevel
    SLam _              -> lamLevel
    SPi _ _             -> piLevel
    SSg _ _             -> sgLevel
    SIf _ _ _           -> ifLevel
    SApp _ _            -> appLevel
    SAnn _ _            -> annLevel

    SThere _            -> initialLevel
    SEnd   _            -> initialLevel
    SRec   _ _          -> initialLevel
    SInit  _            -> initialLevel
    SArg   _ _          -> initialLevel

    SRefl               -> atomicLevel
    SHere               -> atomicLevel
    SWildCard           -> atomicLevel
    SVar _              -> atomicLevel
    SQuotes _           -> atomicLevel

  assoc s = case s of
    SPi   _ _            -> piAssoc
    SSg   _ _            -> sgAssoc
    SApp  _ _            -> appAssoc
    SPair _ _            -> pairAssoc
    SEq   _ _            -> AssocNone

    SInit       _        -> AssocNone
    SRec        _ _      -> AssocNone
    SArg        _ _      -> AssocNone
    SLam _               -> AssocNone
    SIf _ _ _            -> AssocNone
    SThere _             -> AssocNone
    SEnd   _             -> AssocNone
    SAnn _ _             -> AssocNone
    SRefl                -> AssocNone
    SHere                -> AssocNone
    SWildCard            -> AssocNone
    SVar _               -> AssocNone
    SQuotes _            -> AssocNone

----------------------------------------------------------------------

instance Precedence Nom where
  level _ = atomicLevel
  assoc _ = AssocNone

----------------------------------------------------------------------

-- The 'Reader' data of the 'DisplayM'.  E.g., the "flags" Tim
-- mentioned would go in here.
data DisplayR = DisplayR { names :: [Nom] }
emptyDisplayR :: DisplayR
emptyDisplayR = DisplayR { names = [] }

pushName :: Nom -> DisplayM a -> DisplayM a
pushName nm = local (\d -> d { names = nm : names d })

type DisplayM = ReaderT DisplayR FreshM

-- Convert 't' to 'Doc', using 'DisplayR'.
class Display t where
  display :: t -> DisplayM Doc

-- Same as 'Text.Parsec.Expr.Assoc', except with 'Eq' :P
data Assoc = AssocLeft | AssocRight | AssocNone deriving Eq

-- The precedence level is the tightness of binding: larger levels
-- mean tighter binding.
class Precedence t where
  level :: t -> Float
  assoc :: t -> Assoc

-- Precedence levels and infix op associativities.
--
-- Compare with table in used in Spire.Surface.Parsing.
atomicLevel    = -1
initialLevel   = 0

annLevel       = atomicLevel
appLevel       = initialLevel
appAssoc       = AssocLeft
consLevel      = 2
consAssoc      = AssocRight
pairLevel      = 3
pairAssoc      = AssocRight
sgLevel        = 4
sgAssoc        = AssocRight
piLevel        = 5
piAssoc        = AssocRight
ifLevel        = 6
lamLevel       = 7
eqLevel        = 8
defsLevel      = 9
defLevel       = 10

-- For non-infix operators, use 'wrap' to wrap optionally recursive
-- displays as needed.
--
-- For infix operators, use 'wrap' to optionally wrap recursive
-- displays in positions consistent with association, and use
-- 'wrapNonAssoc' to optionally wrap recursive displays in positions
-- inconsistent with association.
--
-- For example, for a left associative application 'App', the
-- 'display' case could be written:
--
--   display s@(App f x) =
--     sepM [ wrap s f (display f) , wrapNonAssoc s x (display x) ]
wrap , wrapNonAssoc :: (Display t2, Precedence t1 , Precedence t2) =>
                            t1 -> t2 -> DisplayM Doc
wrap outside inside =
  if level outside  < level inside ||
     level outside == level inside && assoc outside /= assoc inside
  then parensM . display $ inside
  else display inside

wrapNonAssoc outside inside =
  if level outside <= level inside
  then parensM . display $ inside
  else display inside

----------------------------------------------------------------------