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

import Control.Arrow (Arrow(first))

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

import Presyntax.Presyntax (Plicity(..))

import Prettyprinter.Render.Terminal
import Prettyprinter

import Syntax

instance Pretty Name where
  pretty = pretty . getNameText

prettyTm :: Term -> Doc AnsiStyle
prettyTm = prettyTm . everywhere (mkT beautify) where
  prettyTm (Ref v) =
    case T.uncons (getNameText v) of
      Just ('.', w) -> keyword (pretty w)
      _ -> pretty v
  prettyTm (Con v) = keyword (pretty v)
  prettyTm (PCon _ v) = keyword (pretty v)
  prettyTm (Data _ v) = operator (pretty v)

  prettyTm (App Im f _) = prettyTm f
  prettyTm (App Ex f x) = parenFun f <+> parenArg x

  prettyTm (Pair x y) = parens $ prettyTm x <> operator comma <+> prettyTm y
  prettyTm (Proj1 x) = prettyTm x <> operator (pretty ".1")
  prettyTm (Proj2 x) = prettyTm x <> operator (pretty ".2")

  prettyTm l@(Lam _ _ _) = pretty '\\' <> hsep (map prettyArgList al) <+> pretty "->" <+> nest 2 (prettyTm bod) where
    unwindLam (Lam p x y) = first ((p, x):) (unwindLam y)
    unwindLam (PathIntro _ _ _ (Lam p x y)) = first ((p, x):) (unwindLam y)
    unwindLam t = ([], t)

    (al, bod) = unwindLam l

    used = freeVars bod

    prettyArgList (Ex, v)
      | v `Set.member` used = pretty v
      | otherwise = pretty "_"
    prettyArgList (Im, v)
      | v `Set.member` used = braces $ pretty v
      | otherwise = pretty "_"

  prettyTm (Meta x) = keyword $ pretty '?' <> pretty (mvName x)
  prettyTm Type{}  = keyword $ pretty "Type"
  prettyTm Typeω{} = keyword $ pretty "Typeω"
  prettyTm I{}     = keyword $ pretty "I"
  prettyTm I0{}    = keyword $ pretty "i0"
  prettyTm I1{}    = keyword $ pretty "i1"

  prettyTm (Pi Ex (T.unpack . getNameText -> "_") d r) = prettyDom d <+> pretty "->" <+> prettyTm r
  prettyTm (Pi Im v d r)                 = group (braces (pretty v <+> colon <+> prettyTm d)) <> softline <> pretty "->" <+> prettyTm r
  prettyTm (Pi Ex v d r)                 = group (parens (pretty v <+> colon <+> prettyTm d)) <> softline <> pretty "->" <+> prettyTm r

  prettyTm (Sigma (T.unpack . getNameText -> "_") d r) = prettyDom d <+> pretty "*" <+> prettyTm r
  prettyTm (Sigma v d r)                 = parens (pretty v <+> colon <+> prettyTm d) <+> pretty "*" <+> prettyTm r

  prettyTm (IAnd x y) = parens $ prettyTm x <+> operator (pretty "&&") <+> prettyTm y
  prettyTm (IOr x y) = parens $ prettyTm x <+> operator (pretty "||") <+> prettyTm y
  prettyTm (INot x) = operator (pretty '~') <> prettyTm x

  prettyTm (System (Map.null -> True)) = braces mempty

  prettyTm (System xs) = braces (line <> indent 2 (vcat (punctuate comma (map go (Map.toList xs)))) <> line) where
    go (f, t) = prettyTm f <+> operator (pretty "=>") <+> prettyTm t

  prettyTm (Case _ x xs) = keyword (pretty "case") <+> prettyTm x <+> keyword (pretty "of") <+> braces (line <> indent 2 (prettyCase xs) <> line)
  prettyTm (Let xs e) = align $ keyword (pretty "let") <+> braces (line <> indent 2 (prettyLet xs) <> line) <+> keyword (pretty "in") <+> prettyTm e

  prettyTm x = error (show x)

  prettyCase = vcat . map go where
    go (x, xs) = prettyTm x <+> operator (pretty "=>") <+> prettyTm xs

  prettyLet = vcat . map go where
    go (x, t, y) = pretty x <+> align (colon <+> nest (- 1) (prettyTm t)) <> line <> pretty x <+> pretty "=" <+> prettyTm y

  beautify (PathP l x y)       = toFun "PathP" [l, x, y]
  beautify (IElim _ _ _ f i)   = App Ex f i
  beautify (PathIntro _ _ _ f) = f

  beautify (App _ (Lam _ v b) _)
    | v `Set.notMember` freeVars b = beautify b

  beautify (IsOne phi)  = toFun "IsOne" [phi]
  beautify ItIsOne      = Ref (Bound (T.pack ".1=1") 0)

  beautify (Lam Ex v (App Ex f (Ref v')))
    | v == v', v `Set.notMember` freeVars f = f

  beautify (Comp a I0 (System (Map.null -> True)) a0) = toFun "transp" [a, a0]
  beautify (Lam _ _ (Lam _ _ (System (Map.null -> True)))) = System mempty

  beautify (Partial phi a)   = toFun "Partial" [phi, a]
  beautify (PartialP phi a)  = toFun "PartialP" [phi, a]
  beautify (Comp a phi u a0) = toFun "comp" [a, phi, u, a0]
  beautify (HComp a phi u a0) = toFun "hcomp" [a, phi, u, a0]
  beautify (Sub a phi u)     = toFun "Sub" [a, phi, u]
  beautify (Inc _ _ u)       = toFun "inS" [u]
  beautify (Ouc a phi u u0)  = toFun "outS" [a, phi, u, u0]

  beautify (GlueTy a I1 _ _)    = a
  beautify (GlueTy a b c d)     = toFun "Glue"   [a,b,c,d]
  beautify (Glue a b c d e f)   = toFun "glue"   [a,b,c,d,e,f]
  beautify (Unglue a b c d e)   = toFun "unglue" [a,b,c,d,e]
  beautify x = x

  toFun s a = foldl (App Ex) (Ref (Bound (T.pack ('.':s)) 0)) a

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

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

parenArg :: Term -> Doc AnsiStyle
parenArg x@App{} = parens (prettyTm x)
parenArg x@IElim{} = parens (prettyTm x)
parenArg x@IsOne{} = parens $ prettyTm x
parenArg x@Partial{} = parens $ prettyTm x
parenArg x@PartialP{} = parens $ prettyTm x

parenArg x@Sub{} = parens $ prettyTm x
parenArg x@Inc{} = parens $ prettyTm x
parenArg x@Ouc{} = parens $ prettyTm x

parenArg x@Comp{} = parens $ prettyTm x
parenArg x@Case{} = parens $ prettyTm x

parenArg x = prettyDom x

prettyDom :: Term -> Doc AnsiStyle
prettyDom x@Pi{}    = parens (prettyTm x)
prettyDom x@Sigma{} = parens (prettyTm x)
prettyDom x         = parenFun x

parenFun :: Term -> Doc AnsiStyle
parenFun x@Lam{} = parens $ prettyTm x
parenFun x@PathIntro{} = parens $ prettyTm x
parenFun x = prettyTm x

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

showValue :: Value -> String
showValue = L.unpack . renderLazy . layoutSmart defaultLayoutOptions . prettyTm . quote

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")

freeVars :: Term -> Set Name
freeVars (Ref v) = Set.singleton v
freeVars (App _ f x) = Set.union (freeVars f) (freeVars x)
freeVars (Pi _ n x y) = Set.union (freeVars x) (n `Set.delete` freeVars y)
freeVars (Lam _ n x) = n `Set.delete` freeVars x
freeVars (Let ns x) = Set.union (freeVars x `Set.difference` bound) freed where
  bound = Set.fromList (map (\(x, _, _) -> x) ns)
  freed = foldr (\(_, x, y) -> Set.union (Set.union (freeVars x) (freeVars y))) mempty ns
freeVars Meta{} = mempty
freeVars Con{} = mempty
freeVars PCon{} = mempty
freeVars Data{} = mempty
freeVars Type{} = mempty
freeVars Typeω{} = mempty
freeVars I{} = mempty
freeVars I0{} = mempty
freeVars I1{} = mempty
freeVars (Sigma n x y) = Set.union (freeVars x) (n `Set.delete` freeVars y)
freeVars (Pair x y) = Set.unions $ map freeVars [x, y]
freeVars (Proj1 x) = Set.unions $ map freeVars [x]
freeVars (Proj2 x) = Set.unions $ map freeVars [x]
freeVars (IAnd x y) = Set.unions $ map freeVars [x, y]
freeVars (IOr x y) = Set.unions $ map freeVars [x, y]
freeVars (INot x) = Set.unions $ map freeVars [x]
freeVars (PathP x y z) = Set.unions $ map freeVars [x, y, z]
freeVars (IElim x y z a b) = Set.unions $ map freeVars [x, y, z, a, b]
freeVars (PathIntro x y z a) = Set.unions $ map freeVars [x, y, z, a]
freeVars (IsOne a) = Set.unions $ map freeVars [a]
freeVars ItIsOne{} = mempty
freeVars (Partial x y) = Set.unions $ map freeVars [x, y]
freeVars (PartialP x y) = Set.unions $ map freeVars [x, y]
freeVars (System fs) = foldr (\(x, y) -> Set.union (Set.union (freeVars x) (freeVars y))) mempty (Map.toList fs)

freeVars (Sub x y z) = Set.unions $ map freeVars [x, y, z]
freeVars (Inc x y z) = Set.unions $ map freeVars [x, y, z]
freeVars (Ouc x y z a) = Set.unions $ map freeVars [x, y, z, a]

freeVars (Comp x y z a) = Set.unions $ map freeVars [x, y, z, a]
freeVars (HComp x y z a) = Set.unions $ map freeVars [x, y, z, a]

freeVars (GlueTy x y z a) = Set.unions $ map freeVars [x, y, z, a]
freeVars (Glue x y z a b c) = Set.unions $ map freeVars [x, y, z, a, b, c]
freeVars (Unglue x y z a c) = Set.unions $ map freeVars [x, y, z, a, c]
freeVars (Case rng x y) = freeVars rng <> freeVars x <> foldMap (freeVars . snd) y