module Obs.Typing

import Control.Monad.Maybe

import Data.Nat

import Decidable.Equality

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

import System

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

%default total

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

mismatch : (expected, got : Doc ann) -> Logging ann a
mismatch lhs rhs = fatal $
  hang 2 (pretty "expected" <+> line <+> lhs) <+> comma <+> line <+>
  hang 2 (pretty "got" <+> line <+> align rhs)

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

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

-- Utilities -------------------------------------------------------------------

MkPair : (s, s' : Sort)
  -> NormalForm (isSet s) ctx
  -> NormalForm (isSet s') ctx
  -> NormalForm (isSet (max s s')) ctx
MkPair Prop Prop t u = u
MkPair Prop s'@(Set _) t u = Cnstr $ Pair Prop s' Oh t u
MkPair s@(Set _) s' t u = Cnstr $ Pair s s' Oh t u

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

check : (tyCtx : TyContext ctx)
  -> (s : Sort)
  -> NormalForm True ctx
  -> WithBounds (Term (length ctx))
  -> Logging ann (NormalForm (isSet s) ctx)

checkType : (tyCtx : TyContext ctx)
  -> (s : Sort)
  -> WithBounds (Term (length ctx))
  -> Logging ann (NormalForm True ctx)
checkType ctx s tm = check ctx (suc s) (cast s) tm

infer : (tyCtx : TyContext ctx)
  -> WithBounds (Term (length ctx))
  -> Logging ann (s ** (NormalForm True ctx, NormalForm (isSet s) ctx))

inferType : {default typeMismatch tag : forall a . Doc ann -> Doc ann -> Logging ann a}
  -> (tyCtx : TyContext ctx)
  -> WithBounds (Term (length ctx))
  -> Logging ann (Sort, NormalForm True ctx)
inferType ctx tm = inBounds tm.bounds $ do
  (Set _ ** (Cnstr (Sort s), a)) <- infer ctx tm
    | (_ ** (a, t)) => tag (pretty "sort") (pretty a)
  pure (s, a)

check ctx s ty tm = inBounds tm.bounds $ case (ty, tm.val) of
  (Cnstr (Pi l l' _ a b), Lambda var t) => do
    info "encountered lambda"
    trace $ pretty {ann} "checking body has type" <++> pretty b
    t <- check (ctx ::< MkFreeVar l a var.val) l' b (assert_smaller tm t)
    let Yes Refl = decEq s (l ~> l')
      | No _ => fatal "internal sort mismatch"
    pure $ case l' of
      Prop => Irrel
      Set _ => Cnstr (Lambda l var.val t)
  (_, Lambda var t) => typeMismatch (pretty "function") (pretty ty)
  (Cnstr (Sigma l l' var a b), Pair t u) => do
    info "encountered pair"
    trace $ pretty {ann} "checking first has type" <++> pretty a
    t <- check ctx l a (assert_smaller tm t)
    b <- subst1 t b
    trace $ pretty {ann} "checking second has type" <++> pretty b
    u <- check ctx l' b (assert_smaller tm u)
    let Yes Refl = decEq s (max l l')
      | No _ => fatal "internal sort mismatch"
    pure $ MkPair l l' t u
  (_, Pair t u) => typeMismatch (pretty "pair type") (pretty ty)
  (_, Absurd t) => do
    info "encountered absurd"
    trace "checking proof of contradiction"
    Irrel <- check ctx Prop (Cnstr Bottom) (assert_smaller tm t)
    pure $ doAbsurd _
  (a, t) => do
    info "falling through to inference"
    (l ** (b, t)) <- infer ctx tm
    trace $ pretty {ann} "inferred type is" <++> pretty b
    let Yes Refl = decEq s l
      | No _ => typeMismatch (pretty a) (pretty b)
    a' <- subst a (reduce ctx)
    b' <- subst b (reduce ctx)
    Just _ <- inScope "convert" $ runMaybeT $ convert (suc s) (cast s) a' b'
    | Nothing => inScope "conversion failed" $ mismatch (pretty a) (pretty b)
    pure t

infer ctx tm = inBounds tm.bounds $ case tm.val of
  (Var var i) => do
    info $ pretty {ann} "encountered variable" <++> pretty tm.val
    let def = index' ctx i
    pure (def.sort ** (def.ty, def.tm))
  (Sort s) => do
    info $ pretty {ann} "encountered sort" <++> pretty tm.val
    pure (suc (suc s) ** (cast (suc s), cast s))
  (Pi var a b) => do
    info "encountered pi"
    (s, a) <- inferType ctx (assert_smaller tm a)
    trace $ pretty {ann} "domain type is" <++> pretty a
    (s', b) <- inferType (ctx ::< MkFreeVar s a var.val) (assert_smaller tm b)
    trace $ pretty {ann} "codomain type is" <++> pretty b
    pure (suc (s ~> s') ** (cast (s ~> s'), Cnstr (Pi s s' var.val a b)))
  (Lambda var t) => fatal "cannot infer type of lambda"
  (App t u) => do
    info "encountered application"
    (s ** (ty@(Cnstr (Pi l l' _ a b)), t)) <- infer ctx (assert_smaller tm t)
      | (_ ** (ty, _)) => inBounds t.bounds $ typeMismatch (pretty "function") (pretty ty)
    trace $ pretty {ann} "function has type" <++> pretty ty
    trace $ pretty {ann} "checking argument has type" <++> pretty a
    u <- check ctx l a (assert_smaller tm u)
    trace $ pretty {ann} "argument is well typed"
    let Yes Refl = decEq s (l ~> l')
      | No _ => fatal "internal sort mismatch"
    ty <- subst1 u b
    tm <- doApp t u
    pure (l' ** (ty, rewrite sym $ impredicative l l' in tm))
  (Sigma var a b) => do
    info "encountered sigma"
    (s, a) <- inferType ctx (assert_smaller tm a)
    trace $ pretty {ann} "first type is" <++> pretty a
    (s', b) <- inferType (ctx ::< MkFreeVar s a var.val) (assert_smaller tm b)
    trace $ pretty {ann} "second type is" <++> pretty b
    pure (suc (max s s') ** (cast (max s s'), Cnstr (Sigma s s' var.val a b)))
  (Pair t u) => fatal "cannot infer type of pair"
  (Fst t) => do
    info "encountered first"
    (s ** (ty@(Cnstr (Sigma l l' _ a b)), t)) <- infer ctx (assert_smaller tm t)
      | (_ ** (ty, _)) => inBounds t.bounds $ typeMismatch (pretty "pair type") (pretty ty)
    trace $ pretty {ann} "pair has type" <++> pretty ty
    let Yes Refl = decEq s (max l l')
      | No _ => fatal "internal sort mismatch"
    let ty = a
    tm <- doFst (isSet l) (isSet l') (rewrite sym $ maxIsSet l l' in t)
    pure (l ** (ty, tm))
  (Snd t) => do
    info "encountered second"
    (s ** (ty@(Cnstr (Sigma l l' _ a b)), t)) <- infer ctx (assert_smaller tm t)
      | (_ ** (ty, _)) => inBounds t.bounds $ typeMismatch (pretty "pair type") (pretty ty)
    trace $ pretty {ann} "pair has type" <++> pretty ty
    let Yes Refl = decEq s (max l l')
      | No _ => fatal "internal sort mismatch"
    let t = rewrite sym $ maxIsSet l l' in t
    fst <- doFst (isSet l) (isSet l') t
    ty <- subst1 fst b
    tm <- doSnd (isSet l) (isSet l') t
    pure (l' ** (ty, tm))
  Bool => do
    info "encountered bool"
    pure (Set 1 ** (cast (Set 0), Cnstr Bool))
  True => do
    info "encountered true"
    pure (Set 0 ** (Cnstr Bool, Cnstr True))
  False => do
    info "encountered false"
    pure (Set 0 ** (Cnstr Bool, Cnstr False))
  (If var a t u v) => do
    info "encountered if"
    trace "checking discriminant is bool"
    t <- check ctx (Set 0) (Cnstr Bool) (assert_smaller tm t)
    trace "discriminant is well typed"
    (s, a) <- inferType (ctx ::< MkFreeVar (Set 0) (Cnstr Bool) var.val) (assert_smaller tm a)
    ty <- subst1 (Cnstr True) a
    trace $ pretty {ann} "checking true branch has type" <++> pretty ty
    u <- check ctx s ty (assert_smaller tm u)
    ty <- subst1 (Cnstr False) a
    trace $ pretty {ann} "checking false branch has type" <++> pretty ty
    v <- check ctx s ty (assert_smaller tm v)
    ty <- subst1 t a
    tm <- doIf t u v
    pure (s ** (ty, tm))
  Top => do
    info "encountered top"
    pure (Set 0 ** (cast Prop, Cnstr Top))
  Point => do
    info "encountered point"
    pure (Prop ** (Cnstr Top, Irrel))
  Bottom => do
    info "encountered bottom"
    pure (Set 0 ** (cast Prop, Cnstr Bottom))
  (Absurd t) => fatal "cannot infer type of absurd"
  (Equal t u) => do
    info "encountered equal"
    (s ** (a, t)) <- infer ctx (assert_smaller tm t)
    trace $ pretty {ann} "left side has type" <++> pretty a
    u <- check ctx s a (assert_smaller tm u)
    trace "right side is well typed"
    eq <- doEqual (isSet s) a t u
    pure (Set 0 ** (cast Prop, eq))
  (Refl t) => do
    info "encountered refl"
    (s ** (a, t)) <- infer ctx (assert_smaller tm t)
    trace "argument is well typed"
    eq <- doEqual (isSet s) a t t
    pure (Prop ** (eq, Irrel))
  (Transp b t u t' e) => do
    info "encountered transp"
    (s ** (a, t)) <- infer ctx (assert_smaller tm t)
    trace $ pretty {ann} "index type is" <++> pretty a
    t' <- check ctx s a (assert_smaller tm t')
    trace "new index is well typed"
    b <- check ctx (s ~> Set 0) (Cnstr $ Pi s (Set 0) "" a (cast Prop)) (assert_smaller tm b)
    trace "transporting in Prop"
    oldB <- doApp b t
    u <- check ctx Prop oldB (assert_smaller tm u)
    trace "old-indexed proposition is well typed"
    eq <- doEqual (isSet s) a t t'
    trace $ pretty {ann} "checking equality has type" <++> pretty eq
    e <- check ctx Prop eq (assert_smaller tm e)
    trace "equality is well typed"
    newB <- doApp b t'
    pure (Prop ** (newB, Irrel))
  (Cast a b e t) => do
    info "encountered cast"
    (s, a) <- inferType ctx (assert_smaller tm a)
    trace $ pretty {ann} "old type is" <++> pretty a
    b <- checkType ctx s (assert_smaller tm b)
    trace $ pretty {ann} "new type is" <++> pretty b
    eq <- doEqual True (cast s) a b
    e <- check ctx Prop eq (assert_smaller tm e)
    trace "equality is well typed"
    t <- check ctx s a (assert_smaller tm t)
    trace "value is well typed"
    tm <- doCast (isSet s) a b t
    pure (s ** (b, tm))
  (CastId t) => do
    info "encountered castRefl"
    (s ** (a, t)) <- infer ctx (assert_smaller tm t)
    trace $ pretty {ann} "argument has type" <++> pretty a
    lhs <- doCast (isSet s) a a t
    ret <- doEqual (isSet s) a lhs t
    pure (Prop ** (ret, Irrel))

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

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

export
checkBlock : Block n -> Logging ann (ctx ** (Context ctx, length ctx = n))
checkBlock [] = pure ([] ** ([], Refl))
checkBlock (blk :< def) = do
  (_ ** (ctx, Refl)) <- checkBlock blk
  def <- checkDef ctx def
  pure (_ ** (ctx :< def, Refl))