module Obs.Typing

import Data.Fin

import Obs.NormalForm
import Obs.Sort
import Obs.Substitution
import Obs.Term

import System

import Text.Bounded
import Text.PrettyPrint.Prettyprinter
import Text.PrettyPrint.Prettyprinter.Render.Terminal

%default total

-- Conversion ------------------------------------------------------------------

-- invariant: all definitions fully expanded
-- invariant: m |- t, u <- a : s
convert : (t, u, a : NormalForm n) -> Sort -> Bool
convert Irrel Irrel a Prop = True
convert t u (Ntrl _) (Set k) = t == u
convert (Cnstr Top) (Cnstr Top) (Cnstr (Sort Prop)) (Set 0) = True
convert (Ntrl t) u (Cnstr (Sort s)) (Set k) = Ntrl t == u
convert t (Ntrl u) (Cnstr (Sort s)) (Set k) = t == Ntrl u
convert (Cnstr (Sort s')) (Cnstr (Sort s'')) (Cnstr (Sort s)) (Set k) = s' == s''
convert _ _ _ _ = False

-- Typing Contexts -------------------------------------------------------------

weaken : (m : _) -> NormalForm n -> NormalForm (m + n)
weaken m t = rename t (shift m)

index : Context n -> Fin n -> (NormalForm n, NormalForm n, Sort)
index (ctx :< def) FZ     = bimap (weaken 1) (mapFst (weaken 1)) (def.tm, def.ty, def.sort)
index (ctx :< def) (FS i) = bimap (weaken 1) (mapFst (weaken 1)) $ index ctx i

asSubst : Context n -> Fin n -> NormalForm 0
asSubst (ctx :< def) FZ     = subst def.tm (asSubst ctx)
asSubst (ctx :< def) (FS i) = asSubst ctx i

-- Checking and Inference ------------------------------------------------------

export
Error : Type -> Type -> Type
Error ann = Either (Bounds, Doc ann)

export
printErr : Error ? a -> IO a
printErr (Left (b, p)) = do
  () <- putDoc (pretty "typing: \{show b.startLine}:\{show b.startCol}-\{show b.endLine}:\{show b.endCol}:" <+> softline <+> p)
  exitFailure
printErr (Right x) = pure x

report : Bounds -> Doc ann -> Error ann a
report = curry Left
  
check : Context n -> Term n -> NormalForm n -> Sort -> Error ann (NormalForm n)
infer : Context n -> Term n -> Error ann (NormalForm n, NormalForm n, Sort)

check ctx tm ty s = do
  let bounds = fullBounds tm
  (tm, ty', s') <- infer ctx tm
  let True = s == s'
    | False => report bounds
      (pretty "sort mismatch: expected" <++> pretty s <+> comma <+> softline <+> "got" <++> pretty s')
  let True = convert (subst ty $ asSubst ctx) (subst ty' $ asSubst ctx) (Cnstr (Sort s)) (suc s)
    | False => report bounds
      (pretty "type mismatch: expected" <++> pretty ty <+> comma <+> softline <+> "got" <++> pretty ty')
  pure tm

infer ctx (Var b i) = pure $ index ctx i
infer ctx (Sort b s) = pure $ (Cnstr (Sort s), Cnstr (Sort (suc s)), suc (suc s))
infer ctx (Top b) = pure $ (Cnstr Top, Cnstr (Sort Prop), Set 0)
infer ctx (Point b) = pure $ (Irrel, Cnstr Top, Prop)

-- Checking Definitions and Blocks ---------------------------------------------

checkDef : Context n -> Term.Definition n -> Error ann (NormalForm.Definition n)
checkDef ctx def = do
  let tyBounds = fullBounds def.ty
  let tmBounds = fullBounds def.tm
  (ty, Cnstr (Sort sort), _) <- infer ctx def.ty
    | (_, a, _) => report tyBounds
      (pretty "invalid declaration: expected sort" <+> comma <+> softline <+> "got" <++> pretty a)
  tm <- check ctx def.tm ty sort
  pure $ MkDefinition {name = Just def.name, ty, tm, sort}

export
checkBlock : Block n -> Error ann (Context n)
checkBlock [] = pure []
checkBlock (blk :< def) = do
  ctx <- checkBlock blk
  def <- checkDef ctx def
  pure (ctx :< def)