||| Raw syntax for terms. This is before eliminating names.
module Obs.Syntax

import Obs.Universe

import Text.Bounded
import Text.PrettyPrint.Prettyprinter

%default total

-- Definition ------------------------------------------------------------------

public export
data Syntax : Type where
  Var : (var : String) -> Syntax
  -- Universes
  Universe : (s : Universe) -> Syntax
  -- Dependent Products
  Pi : (var : WithBounds String) -> (domain, codomain : WithBounds Syntax) -> Syntax
  Lambda : (var : WithBounds String) -> (body : WithBounds Syntax) -> Syntax
  App : (fun, arg : WithBounds Syntax) -> Syntax
  -- Dependent Sums
  Sigma : (var : WithBounds String) -> (index, element : WithBounds Syntax) -> Syntax
  Pair : (first, second : WithBounds Syntax) -> Syntax
  First : (arg : WithBounds Syntax) -> Syntax
  Second : (arg : WithBounds Syntax) -> Syntax
  -- Booleans
  Bool : Syntax
  True : Syntax
  False : Syntax
  If : (var : WithBounds String) ->
    (returnType, discriminant, true, false : WithBounds Syntax) ->
    Syntax
  -- True
  Top : Syntax
  Point : Syntax
  -- False
  Bottom : Syntax
  Absurd : (contradiction : WithBounds Syntax) -> Syntax
  -- Equality
  Equal : (left, right : WithBounds Syntax) -> Syntax
  Refl : (value : WithBounds Syntax) -> Syntax
  Transp : (indexedType, oldIndex, value, newIndex, equality : WithBounds Syntax) -> Syntax
  Cast : (oldType, newType, equality, value : WithBounds Syntax) -> Syntax
  CastId : (value : WithBounds Syntax) -> Syntax

public export
record Definition where
  constructor MkDefinition
  name   : WithBounds String
  ty     : WithBounds Syntax
  tm     : WithBounds Syntax

-- Pretty Print ----------------------------------------------------------------

prettyDecl : (name : String) -> (ty : Doc ann) -> Doc ann
prettyDecl name ty = parens $ group (pretty name <++> colon <+> line <+> align ty)

prettyDeclForm : (prec : Prec) ->
  (name : String) ->
  (ty : Doc ann) ->
  (sep : String) ->
  (tail : Doc ann) ->
  Doc ann
prettyDeclForm d name ty sep tail =
  parenthesise (d > Open) $
  group $
  vsep $
  [prettyDecl name ty <++> pretty sep, align tail]

prettyLambda : (prec : Prec) -> (name : String) -> (body : Doc ann) -> Doc ann
prettyLambda d name body =
  parenthesise (d >= App) $
  group $
  backslash <+> pretty name <++> pretty "=>" <+> line <+> align body

prettyApp : (prec : Prec) -> (head : Doc ann) -> (tail : List (Doc ann)) -> Doc ann
prettyApp d head tail = parenthesise (d >= App) $ group $ vsep (align head :: map (hang 2) tail)

prettyPrec : Prec -> Syntax -> Doc ann
prettyPrecBounds : Prec -> WithBounds Syntax -> Doc ann

prettyPrec d (Var {var}) = pretty var
prettyPrec d (Universe {s}) = prettyPrec d s
prettyPrec d (Pi {var, domain, codomain}) =
  prettyDeclForm d var.val (prettyPrecBounds Open domain) "->" (prettyPrecBounds Open codomain)
prettyPrec d (Lambda {var, body}) = prettyLambda d var.val (prettyPrecBounds Open body)
prettyPrec d (App {fun, arg}) = prettyApp d (prettyPrecBounds Open fun) [prettyPrecBounds App arg]
prettyPrec d (Sigma {var, index, element}) =
  prettyDeclForm d var.val (prettyPrecBounds Open index) "**" (prettyPrecBounds Open element)
prettyPrec d (Pair {first, second}) =
  angles $
  group $
  neutral <++>
  align (prettyPrecBounds Open first) <+>
  comma <+>
  line <+>
  align (prettyPrecBounds Open second) <++>
  neutral
prettyPrec d (First {arg}) = prettyApp d (pretty "fst") [prettyPrecBounds App arg]
prettyPrec d (Second {arg}) = prettyApp d (pretty "snd") [prettyPrecBounds App arg]
prettyPrec d Bool = pretty "Bool"
prettyPrec d True = pretty "True"
prettyPrec d False = pretty "False"
prettyPrec d (If {var, returnType, discriminant, true, false}) =
  prettyApp d (pretty "if") $
  [ prettyLambda App var.val (prettyPrecBounds Open returnType)
  , prettyPrecBounds App discriminant
  , prettyPrecBounds App true
  , prettyPrecBounds App false
  ]
prettyPrec d Top = pretty "()"
prettyPrec d Point = pretty "tt"
prettyPrec d Bottom = pretty "Void"
prettyPrec d (Absurd {contradiction}) =
  prettyApp d (pretty "absurd") [prettyPrecBounds App contradiction]
prettyPrec d (Equal {left, right}) =
  parenthesise (d >= Equal) $
  group $
  align (prettyPrecBounds Equal left) <++>
  pretty "~" <+>
  line <+>
  align (prettyPrecBounds Equal right)
prettyPrec d (Refl {value}) = prettyApp d (pretty "refl") [prettyPrecBounds App value]
prettyPrec d (Transp {indexedType, oldIndex, value, newIndex, equality}) =
  prettyApp d (pretty "transp") $
  [ prettyPrecBounds App indexedType
  , prettyPrecBounds App oldIndex
  , prettyPrecBounds App value
  , prettyPrecBounds App newIndex
  , prettyPrecBounds App equality
  ]
prettyPrec d (Cast {oldType, newType, equality, value}) =
  prettyApp d (pretty "cast") $
  [ prettyPrecBounds App oldType
  , prettyPrecBounds App newType
  , prettyPrecBounds App equality
  , prettyPrecBounds App value
  ]
prettyPrec d (CastId {value}) = prettyApp d (pretty "castRefl") [prettyPrecBounds App value]

prettyPrecBounds d (MkBounded val _ _) = prettyPrec d val

export
Pretty Syntax where
  prettyPrec = Syntax.prettyPrec

export
Pretty Definition where
  pretty def = group $
    pretty def.name.val <++> colon <+> softline <+> pretty def.ty.val <+> hardline <+>
    pretty def.name.val <++> equals <+> softline <+> pretty def.tm.val