||| Abstract syntax for terms.
module Obs.Term

import Data.Fin

import Obs.Pretty
import Obs.Universe

import Text.Bounded

%default total

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

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

public export
record Definition (n : Nat) where
  constructor MkDefinition
  name : WithBounds String
  ty   : WithBounds (Term n)
  tm   : WithBounds (Term n)

public export
data Block : Nat -> Type where
  Nil  : Block 0
  (:<) : Block n -> Definition n -> Block (S n)

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

prettyPrecBounds : Prec -> WithBounds (Term n) -> Doc ann

prettyPrec : Prec -> Term n -> Doc ann
prettyPrec d (Var {var, i}) = 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}) =
  prettyPair (prettyPrecBounds Open first) (prettyPrecBounds Open second)
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}) =
  prettyEqual d (prettyPrecBounds Equal left) (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 v _ _) = prettyPrec d v

export
Pretty (Term n) where
  prettyPrec = Term.prettyPrec

export
Pretty (Definition n) 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

export
Pretty (Block n) where
  pretty []           = neutral
  pretty ([] :< def)  = pretty def
  pretty (blk :< def) = pretty blk <+> hardline <+> hardline <+> pretty def