module Obs.Typing

import Data.Vect

import Obs.Logging
import Obs.NormalForm
import Obs.NormalForm.Normalise
import Obs.Sort
import Obs.Substitution
import Obs.Term
import Obs.Typing.Conversion

import System

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

%default total

-- Loggings ----------------------------------------------------------------------

Rename (Logging ann . NormalForm) where
  rename t f = pure $ rename !t f

mismatch : Bounds -> Doc ann -> Doc ann -> Logging ann a
mismatch bounds lhs rhs = fatal $
  MkBounded
    (pretty "expected" <++> lhs <+> comma <+> softline <+> pretty "got" <++> rhs)
    False
    bounds

typeMismatch : Bounds -> Doc ann -> Doc ann -> Logging ann a
typeMismatch bounds lhs rhs = inScope "type mismatch" $ mismatch bounds lhs rhs

sortMismatch : Bounds -> Doc ann -> Doc ann -> Logging ann a
sortMismatch bounds lhs rhs = inScope "sort mismatch" $ mismatch bounds lhs rhs

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

infix 5 ::<

data TyContext : Nat -> Nat -> Type where
  Nil   : TyContext 0 0
  (:<)  : TyContext m n -> NormalForm.Definition n -> TyContext m (S n)
  (::<) : TyContext m n -> (String, NormalForm n, Sort) -> TyContext (S m) (S n)

fromContext : Context n -> TyContext 0 n
fromContext []           = []
fromContext (ctx :< def) = fromContext ctx :< def

countVars : TyContext m n -> Fin (S n)
countVars []          = FZ
countVars (ctx :< y)  = weaken $ countVars ctx
countVars (ctx ::< y) = FS $ countVars ctx

index : TyContext m n -> Fin n -> (NormalForm n, NormalForm n, Sort)
index (ctx :< def)              FZ     = (weaken 1 def.tm, weaken 1 def.ty, def.sort)
index (ctx ::< (var, ty, Prop)) FZ     = (Irrel, weaken 1 ty, Prop)
index (ctx ::< (var, ty, s))    FZ     = (Ntrl $ Var var FZ, weaken 1 ty, s)
index (ctx :< _)                (FS i) = bimap (weaken 1) (mapFst (weaken 1)) $ index ctx i
index (ctx ::< _)               (FS i) = bimap (weaken 1) (mapFst (weaken 1)) $ index ctx i

partial
asSubst : (ctx : TyContext m n) -> Fin n -> Logging ann (NormalForm m)
asSubst (ctx :< def)             FZ     = subst def.tm (asSubst ctx)
asSubst (ctx ::< (var, _, Prop)) FZ     = pure Irrel
asSubst (ctx ::< (var, _, _))    FZ     = pure $ Ntrl (Var var FZ)
asSubst (ctx :< def)             (FS i) = asSubst ctx i
asSubst (ctx ::< _)              (FS i) = map (weaken 1) (asSubst ctx i)

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

partial
check : TyContext m n -> WithBounds (Term n) -> NormalForm n -> Sort -> Logging ann (NormalForm n)
partial
infer : TyContext m n -> WithBounds (Term n) -> Logging ann (NormalForm n, NormalForm n, Sort)
partial
inferType : {default typeMismatch tag : forall a . Bounds -> Doc ann -> Doc ann -> Logging ann a}
  -> TyContext m n -> WithBounds (Term n) -> Logging ann (NormalForm n, Sort)

check ctx tm ty s = case (tm.val, ty) of
  (Lambda _ t, Cnstr (Pi s s' var a b)) => do
    inScope "check" $ trace $ map (\_ => pretty {ann} "checking under lambda with type" <++> pretty a) tm
    inScope "check" $ trace $ map (\_ => "binding new variable to \{var}") tm
    inScope "check" $ trace $ map (\_ => pretty {ann} "checking for type" <++> pretty b) tm
    t <- check (ctx ::< (var, a, s)) t b s'
    case s' of
      Prop => pure Irrel
      _ => pure (Cnstr $ Lambda var t)
  (Lambda _ _, ty) => typeMismatch tm.bounds (pretty "pi") (pretty ty)
  (Pair t u, Cnstr (Sigma s s' var a b)) => do
    inScope "check" $ trace $ map (\_ => "checking pair") tm
    inScope "check" $ trace $ map (\_ => pretty {ann} "checking first for type" <++> pretty a) tm
    t <- check ctx t a s
    b <- subst1 t b
    inScope "check" $ trace $ map (\_ => pretty {ann} "checking second for type" <++> pretty b) tm
    u <- check ctx u b s'
    inScope "check" $ trace $ map (\_ => pretty {ann} "pair is well typed") tm
    case (s, s') of
      (Prop, Prop) => pure Irrel
      _ => pure (Cnstr $ Pair t u)
  (Pair _ _, ty) => typeMismatch tm.bounds (pretty "sigma") (pretty ty)
  (Left t, Cnstr (Sum s s' a b)) => do
    inScope "check" $ trace $ map (\_ => "checking left injection") tm
    inScope "check" $ trace $ map (\_ => pretty {ann} "checking subterm for type" <++> pretty a) tm
    t <- check ctx t a s
    inScope "check" $ trace $ map (\_ => " subterm is well typed") tm
    pure (Cnstr $ Left t)
  (Left _, ty) => typeMismatch tm.bounds (pretty "sum") (pretty ty)
  (Right t, Cnstr (Sum s s' a b)) => do
    inScope "check" $ trace $ map (\_ => "checking right injection") tm
    inScope "check" $ trace $ map (\_ => pretty {ann} "checking subterm for type" <++> pretty b) tm
    t <- check ctx t b s'
    inScope "check" $ trace $ map (\_ => " subterm is well typed") tm
    pure (Cnstr $ Right t)
  (Right _, ty) => typeMismatch tm.bounds (pretty "sum") (pretty ty)
  (_, _) => do
    inScope "check" $ trace $ map (\_ => "checking has fallen through") tm
    (v, a, s') <- infer ctx tm
    inScope "check" $ trace $ map (\_ => pretty {ann} "inferred type is" <++> pretty a) tm
    let True = s == s'
      | False => sortMismatch tm.bounds (pretty s) (pretty s')
    True <- convert !(subst ty $ asSubst ctx) !(subst a $ asSubst ctx) (cast s) (suc s)
      | False => typeMismatch tm.bounds (pretty ty) (pretty a)
    inScope "check" $ trace $ map (\_ => pretty {ann} "converted" <++> pretty a <+> softline <+> pretty "to" <++> pretty ty) tm
    pure v

infer ctx tm = case tm.val of
  (Var var i) => do
    inScope "infer" $ trace $ map (\_ => "encountered variable \{var}@\{show i}") tm
    let (t, a, s) = index ctx i
    inScope "infer" $ trace $ map (\_ => pretty {ann} "variable has type" <++> pretty a) tm
    pure (t, a, s)
  (Sort s) => pure (cast s, cast (suc s), suc (suc s))
  (Pi var a b) => do
    inScope "infer" $ trace $ map (\_ => "encountered Pi type") tm
    (a, s) <- inferType ctx a
    inScope "infer" $ trace $ map (\_ => pretty {ann} "argument has type" <++> pretty a) tm
    (b, s') <- inferType (ctx ::< (var.val, a, s)) b
    inScope "infer" $ trace $ map (\_ => pretty {ann} "result has type" <++> pretty b <+> comma <+> softline <+> pretty "so Pi type has type" <++> pretty (s ~> s')) tm
    pure (Cnstr (Pi s s' var.val a b), cast (s ~> s'), suc (s ~> s'))
  (Lambda _ _) => inScope "cannot infer type" $ fatal tm
  (App t u) => do
    inScope "infer" $ trace $ map (\_ => "encountered application") tm
    (t, ty@(Cnstr (Pi s s' _ a b)), _) <- infer ctx t
      | (_, ty, _) => inScope "wrong type to apply" $ fatal (map (\_ => ty) tm)
    inScope "infer" $ trace $ map (\_ => pretty {ann} "function has type" <++> pretty ty) tm
    inScope "infer" $ trace $ map (\_ => pretty {ann} "checking argument has type" <++> pretty a) tm
    u <- check ctx u a s
    inScope "infer" $ trace $ map (\_ => "argument is well typed") tm
    res <- doApp t u
    ty <- subst1 u b
    inScope "infer" $ trace $ map (\_ => pretty {ann} "final result is" <++> pretty res <+> softline <+> pretty "of type" <++> pretty ty) tm
    pure (res, ty, s')
  (Sigma var a b) => do
    inScope "infer" $ trace $ map (\_ => "encountered Sigma type") tm
    (a, s) <- inferType ctx a
    inScope "infer" $ trace $ map (\_ => pretty {ann} "first has type" <++> pretty a) tm
    (b, s') <- inferType (ctx ::< (var.val, a, s)) b
    inScope "infer" $ trace $ map (\_ => pretty {ann} "second has type" <++> pretty b <+> comma <+> softline <+> pretty "so Sigma type has type" <++> pretty (lub s s')) tm
    pure (Cnstr (Sigma s s' var.val a b), cast (lub s s'), suc (lub s s'))
  (Pair _ _) => inScope "cannot infer type" $ fatal tm
  (Fst t) => do
    inScope "infer" $ trace $ map (\_ => "encountered first projection") tm
    (t, ty@(Cnstr (Sigma s s' _ a b)), _) <- infer ctx t
      | (_, ty, _) => inScope "wrong type to fst" $ fatal (map (\_ => ty) tm)
    inScope "infer" $ trace $ map (\_ => pretty {ann} "pair has type" <++> pretty ty) tm
    res <- doFst t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "final result is" <++> pretty res <+> softline <+> pretty "of type" <++> pretty a) tm
    pure (res, a, s)
  (Snd t) => do
    inScope "infer" $ trace $ map (\_ => "encountered first projection") tm
    (t, ty@(Cnstr (Sigma s s' _ a b)), _) <- infer ctx t
      | (_, ty, _) => inScope "wrong type to snd" $ fatal (map (\_ => ty) tm)
    inScope "infer" $ trace $ map (\_ => pretty {ann} "pair has type" <++> pretty ty) tm
    t1 <- doFst t
    res <- doSnd t
    b <- subst1 t1 b
    inScope "infer" $ trace $ map (\_ => pretty {ann} "final result is" <++> pretty res <+> softline <+> pretty "of type" <++> pretty b) tm
    pure (res, b, s')
  (Sum a b) => do
    inScope "infer" $ trace $ map (\_ => "encountered sum type") tm
    (a, s) <- inferType ctx a
    inScope "infer" $ trace $ map (\_ => pretty {ann} "left has type" <++> pretty a) tm
    (b, s') <- inferType ctx b
    inScope "infer" $ trace $ map (\_ => pretty {ann} "right has type" <++> pretty b) tm
    pure (Cnstr (Sum s s' a b), cast (ensureSet (lub s s')), suc (ensureSet (lub s s')))
  (Left _) => inScope "cannot infer type" $ fatal tm
  (Right _) => inScope "cannot infer type" $ fatal tm
  (Case t l out f g) => do
    inScope "infer" $ trace $ map (\_ => "encountered case") tm
    (t, ty@(Cnstr (Sum s s' a b)), s'') <- infer ctx t
      | (_, ty, _) => inScope "wrong type to case" $ fatal (map (\_ => ty) tm)
    inScope "infer" $ trace $ map (\_ => pretty {ann} "sum has type" <++> pretty ty) tm
    (Cnstr (Sort l), _) <- inferType ctx l
      | (s, _) => typeMismatch l.bounds (pretty "sort") (pretty s)
    inScope "infer" $ trace $ map (\_ => pretty {ann} "output has sort" <++> pretty l) tm
    let ty' = Cnstr (Pi s'' (suc l) "x" ty (cast l))
    inScope "infer" $ trace $ map (\_ => pretty {ann} "checking output has type" <++> pretty ty') tm
    out <- check ctx out ty' (s'' ~> suc l)
    inScope "infer" $ trace $ map (\_ => "output is well typed") tm
    outL <- doApp (weaken 1 out) (Cnstr (Left (Ntrl (Var "x" 0))))
    let fTy = Cnstr (Pi s l "x" a outL)
    inScope "infer" $ trace $ map (\_ => pretty {ann} "checking left branch has type" <++> pretty fTy) tm
    f <- check ctx f fTy (s ~> l)
    inScope "infer" $ trace $ map (\_ => "left branch is well typed") tm
    outR <- doApp (weaken 1 out) (Cnstr (Right (Ntrl (Var "x" 0))))
    let gTy = Cnstr (Pi s' l "x" b outR)
    inScope "infer" $ trace $ map (\_ => pretty {ann} "checking right branch has type" <++> pretty gTy) tm
    g <- check ctx g gTy (s' ~> l)
    inScope "infer" $ trace $ map (\_ => "right branch is well typed") tm
    res <- doCase t f g
    out <- doApp out t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "final result is" <++> pretty res <+> softline <+> pretty "of type" <++> pretty out) tm
    pure (res, out, l)
  Top => pure $ (Cnstr Top, cast Prop, Set 0)
  Point => pure $ (Irrel, Cnstr Top, Prop)
  Bottom  => pure $ (Cnstr Bottom, cast Prop, Set 0)
  (Absurd a t) => do
    inScope "infer" $ trace $ map (\_ => "encountered absurd") tm
    (a, s) <- inferType ctx a
    inScope "infer" $ trace $ map (\_ => pretty {ann} "will fudge to type" <++> pretty a) tm
    inScope "infer" $ trace $ map (\_ => pretty {ann} "checking for proof of false") tm
    _ <- check ctx t (Cnstr Bottom) Prop
    inScope "infer" $ trace $ map (\_ => "proof of false is well typed") tm
    pure (Ntrl Absurd, a, s)
  (Equal t u) => do
    inScope "infer" $ trace $ map (\_ => "encountered equal") tm
    (t, a, s) <- infer ctx t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "lhs has type" <++> pretty a) tm
    inScope "infer" $ trace $ map (\_ => "checking rhs has same type") tm
    u <- check ctx u a s
    inScope "infer" $ trace $ map (\_ => "rhs is well typed") tm
    res <- doEqual a t u
    inScope "infer" $ trace $ map (\_ => pretty {ann} "equality computes to" <++> pretty res) tm
    pure (res, cast Prop, Set 0)
  (Refl t) => do
    inScope "infer" $ trace $ map (\_ => "encountered refl") tm
    (t, a, s) <- infer ctx t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "term has type" <++> pretty a) tm
    ty <- doEqual a t t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "equality computes to" <++> pretty ty) tm
    pure (Irrel, ty, Prop)
  (Transp t b u t' e) => do
    inScope "infer" $ trace $ map (\_ => "encountered transp") tm
    (t, a, s) <- infer ctx t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "start index has type" <++> pretty a) tm
    inScope "infer" $ trace $ map (\_ => "checking end index has same type") tm
    t' <- check ctx t' a s
    inScope "infer" $ trace $ map (\_ => "end index is well typed") tm
    let ty = Cnstr (Pi s (Set 0) "_" a (cast Prop))
    inScope "infer" $ trace $ map (\_ => pretty {ann} "checkout output is in" <++> pretty ty) tm
    b <- check ctx b ty s
    inScope "infer" $ trace $ map (\_ => pretty {ann} "output is well typed") tm
    inScope "infer" $ trace $ map (\_ => "checking equality type") tm
    eq <- doEqual a t t'
    _ <- check ctx e eq Prop
    inScope "infer" $ trace $ map (\_ => "equality is well typed") tm
    inScope "infer" $ trace $ map (\_ => "checking transformed value") tm
    inTy <- doApp b t
    _ <- check ctx u inTy Prop
    inScope "infer" $ trace $ map (\_ => "transformed value is well typed") tm
    outTy <- doApp b t'
    inScope "infer" $ trace $ map (\_ => pretty {ann} "producing value of type" <++> pretty outTy) tm
    pure (Irrel, outTy, Prop)
  (Cast b e t) => do
    inScope "infer" $ trace $ map (\_ => "encountered cast") tm
    (t, a, s) <- infer ctx t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "input has type" <++> pretty a) tm
    inScope "infer" $ trace $ map (\_ => "checking output has same sort") tm
    b <- check ctx b (cast s) (suc s)
    inScope "infer" $ trace $ map (\_ => "output is well sorted") tm
    inScope "infer" $ trace $ map (\_ => "checking equality type") tm
    eq <- doEqual (cast s) a b
    _ <- check ctx e eq Prop
    inScope "infer" $ trace $ map (\_ => "equality is well typed") tm
    inScope "infer" $ trace $ map (\_ => pretty {ann} "producing value of type" <++> pretty b) tm
    res <- doCastL a b t
    pure (res, b, s)
  (CastId t) => do
    inScope "infer" $ trace $ map (\_ => "encountered cast identity proof") tm
    (t, a, s) <- infer ctx t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "term has type" <++> pretty a) tm
    cast <- doCastL a a t
    eq <- doEqual a cast t
    inScope "infer" $ trace $ map (\_ => pretty {ann} "producing equality type" <++> pretty eq) tm
    pure (Irrel, eq, Prop)

inferType ctx a = do
  (a, Cnstr (Sort s), _) <- infer ctx a
    | (_, ty, _) => tag a.bounds (pretty "sort") (pretty ty)
  pure (a, s)

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

partial
checkDef : Context n -> Term.Definition n -> Logging ann (NormalForm.Definition n)
checkDef ctx def = do
  (ty, sort) <-
    inferType {tag = \bounds, lhs, rhs => inScope "invalid declaration" $ mismatch bounds lhs rhs}
      (fromContext ctx) def.ty
  inScope "check" $ debug $ map (\name => pretty {ann} "\{name} has type" <++> pretty ty) def.name
  tm <- check (fromContext ctx) def.tm ty sort
  inScope "check" $ debug $ map (\name => pretty {ann} "\{name} is well typed with value" <++> pretty tm) def.name
  pure $ MkDefinition {name = def.name, ty, tm, sort}

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