{-# LANGUAGE LambdaCase #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE NamedFieldPuns #-}
module Syntax.Pretty where


import qualified Data.Map.Strict as Map
import qualified Data.Text.Lazy as L
import qualified Data.Text as T
import Data.Map.Strict (Map)
import Data.Text (Text)

import Presyntax.Presyntax (Plicity(..))

import Prettyprinter.Render.Terminal
import Prettyprinter

import Syntax

instance Pretty Name where
  pretty = pretty . getNameText

prettyTm :: Term -> Doc AnsiStyle
prettyTm = go True 0 where
  go t p =
    \case
      Ref v    -> pretty v
      Con v    -> keyword $ pretty v
      PCon _ v -> keyword $ pretty v
      Data _ v -> keyword $ pretty v

      App Im f _ -> go t p f
      App Ex f x ->
        parenIf (p >= arg_prec) $
              go False fun_prec f
          <+> group (go False arg_prec x)
      
      Lam Ex v (App Ex f (Ref v')) | v == v' -> instead f
      Lam i v t ->
        let
          getArgs (Lam i v t) =
            let (as, b) = getArgs t in ((i, v):as, b)
          getArgs (PathIntro _ _ _ (Lam _ v t)) =
            let (as, b) = getArgs t in ((Ex, v):as, b)
          getArgs t = ([], t)

          (as, b) = getArgs (Lam i v t)
        in
          parenIf (p >= fun_prec) . group $
            pretty '\\' <> hsep (map (\(i, v) -> braceIf (i == Im) (pretty v)) as)
              <+> arrow
              <+> nest 2 (go False 0 b)
      
      Pi _ (T.unpack . getNameText -> "_") d r ->
        parenIf (p >= fun_prec) $
          group (go False dom_prec d)
            <> space <> arrow <> sp
            <> go t 0 r

      Pi i x d r ->
        let
          c = case r of
            Pi _ (getNameText -> x) _ _ | x /= T.pack "_" -> sp
            _ -> space <> arrow <> sp
        in
          parenIf (p >= fun_prec) $
            plic i (pretty x <+> colon <+> go False 0 d)
              <> c <> go t 0 r
      
      Let binds body ->
        parenIf (p >= fun_prec) $
          align $ keyword (pretty "let")
              <> line
              <> indent 2 (prettyBinds binds)
              <> keyword (pretty "in")
              <+> go False 0 body

      Meta MV{mvName} -> keyword (pretty mvName)

      Type -> keyword (pretty "Type")
      Typeω -> keyword (pretty "Pretype")

      Sigma v d r ->
        parenIf (p >= fun_prec) . align $
          group (parens (pretty v <+> colon <+> go False 0 d))
            <+> operator (pretty "*") <+> go False 0 r

      Pair a b -> parens $ go False 0 a <> comma <+> go False 0 b
      Proj1 a -> parenIf (p >= arg_prec) $ go False 0 a <> keyword (pretty ".1")
      Proj2 a -> parenIf (p >= arg_prec) $ go False 0 a <> keyword (pretty ".2")

      I -> keyword (pretty "I")
      I0 -> keyword (pretty "i0")
      I1 -> keyword (pretty "i1")

      IAnd x y -> parenIf (p > and_prec) $
        go False and_prec x <+> operator (pretty "/\\") <+> go False and_prec y

      IOr x y -> parenIf (p > or_prec) $
        go False or_prec x <+> operator (pretty "\\/") <+> go False or_prec y

      INot x -> operator (pretty "~") <> go False p x

      PathP _ x y -> parenIf (p >= arg_prec) $
        go False 0 x <+> operator (pretty "≡") <+> go False 0 y
      
      IElim _a _x _y f i -> instead (App Ex f i)
      PathIntro _a _x _y f -> instead f

      IsOne p -> brackets (go False 0 p)
      ItIsOne -> keyword (pretty "1=1")

      Partial a p -> apps (con "Partial") [(Ex, a), (Ex, p)]
      PartialP a p -> apps (con "PartialP") [(Ex, a), (Ex, p)]

      System fs | Map.null fs -> braces mempty
      System fs ->
        let
          face (f, t) = go False 0 f <+> operator (pretty "=>") <+> align (go False 0 t)
        in
          braces (line <> indent 2 (vsep (map face (Map.toList fs))) <> line)

      Sub a phi u -> apps (con "Sub") [(Ex, a), (Ex, phi), (Ex, u)]
      Inc _ _hi u -> apps (con "inS") [(Ex, u)]
      Ouc _ _hi _ a0 -> apps (con "outS") [(Ex, a0)]

      GlueTy a phi t e -> apps (con "primGlue") [(Ex, a), (Ex, phi), (Ex, t), (Ex, e)]
      Glue _a _phi _ty _e t im -> apps (con "glue") [(Ex, t), (Ex, im)]
      Unglue _a _phi _ty _e t -> apps (con "unglue") [(Ex, t)]

      Comp a phi u a0 -> apps (con "comp") [(Ex, a), (Ex, phi), (Ex, u), (Ex, a0)]
      HComp a phi u a0 -> apps (con "hcomp") [(Ex, a), (Ex, phi), (Ex, u), (Ex, a0)]

      Case _ t cs ->
        let
          oneCase (c, 0, l) = go False 0 c <+> operator (pretty "=>") <+> go False 0 l
          oneCase (c, i, l) =
            let (args, bd) = getLams i l
             in go False 0 c <+> hsep (map pretty args) <+> operator (pretty "=>") <+> go False 0 bd

          getLams 0 x = ([], x)
          getLams n (Lam _ v t) = let (as, b) = getLams (n - 1) t in (v:as, b)
          getLams _ x = ([], x)
        in
          parenIf (p >= fun_prec) $
            keyword (pretty "case") <+> go False 0 t <+> keyword (pretty "of")
                <> line
                <> indent 2 (vsep (map oneCase cs))

      EqS _ x y -> parenIf (p >= arg_prec) $
        go False 0 x <+> operator (pretty "≡S") <+> go False 0 y

      Syntax.Refl _ _ -> keyword (pretty "refl")
      Syntax.AxK _ _ bigp pr eq   -> apps (con "K_s") [(Ex, bigp), (Ex, pr), (Ex, eq)]
      Syntax.AxJ _ _ bigp pr _ eq -> apps (con "J_s") [(Ex, bigp), (Ex, pr), (Ex, eq)]
    where
      sp | t         = softline
         | otherwise = space

      parenIf p x | p = parens x
                  | otherwise = x

      braceIf p x | p = braces x
                  | otherwise = x

      con x = Con (Bound (T.pack x) 0)
      
      plic = \case
        Ex -> parens
        Im -> braces

      arrow = operator (pretty "->")

      instead = go t p
      apps :: Term -> [(Plicity, Term)] -> Doc AnsiStyle 
      apps f xs = instead (foldl (\f (p, x) -> App p f x) f xs)

prettyBinds :: [(Name, Term, Term)] -> Doc AnsiStyle
prettyBinds [] = mempty
prettyBinds ((x, ty, tm):bs) =
     pretty x <+> colon <+> align (prettyTm ty)
  <> line
  <> pretty x <+> equals <+> align (prettyTm tm)
  <> line
  <> prettyBinds bs

showFace :: Map Head Bool -> Doc AnsiStyle
showFace = hsep . map go . Map.toList where
  go (h, b) = parens $ prettyTm (quote (VNe h mempty)) <+> operator (pretty "=") <+> pretty (if b then "i1" else "i0")

prettyVl :: Value -> Doc AnsiStyle
prettyVl = prettyTm . quote

render :: Doc AnsiStyle -> Text
render = L.toStrict . renderLazy . layoutSmart defaultLayoutOptions

arg_prec, fun_prec, dom_prec, and_prec, or_prec :: Int
dom_prec = succ fun_prec
arg_prec = succ and_prec
and_prec = succ or_prec
or_prec  = succ fun_prec
fun_prec = 1

keyword :: Doc AnsiStyle -> Doc AnsiStyle
keyword = annotate (color Magenta)

operator :: Doc AnsiStyle -> Doc AnsiStyle
operator = annotate (color Yellow)