module Obs.Typing

import Data.Vect

import Obs.Logging
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

-- 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

-- Substitution ----------------------------------------------------------------

wkn : Vect k String -> (Fin n -> Logging ann (NormalForm m)) -> Fin (k + n) -> Logging ann (NormalForm (k + m))
wkn []            f = f
wkn (var :: vars) f =
  add
    (Logging ann . NormalForm)
    (pure $ Ntrl $ Var var FZ)
    (\i => pure $ rename !(wkn vars f i) FS)

covering partial
substCnstr : Constructor n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (Constructor m)
covering partial
substNtrl  : Neutral n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (NormalForm m)
covering partial
subst      : NormalForm n -> (Fin n -> Logging ann (NormalForm m)) -> Logging ann (NormalForm m)
covering partial
doApp      : NormalForm n -> NormalForm n -> Logging ann (NormalForm n)

substCnstr (Sort s) f = pure $ Sort s
substCnstr (Pi s s' var a b) f = do
  a <- subst a f
  b <- subst b (wkn [var] f)
  pure (Pi s s' var a b)
substCnstr (Lambda var t) f = do
  t <- subst t (wkn [var] f)
  pure (Lambda var t)
substCnstr Top f = pure $ Top
substCnstr Bottom  f = pure $ Bottom

substNtrl (Var var i) f = do
  Ntrl (Var "" j) <- f i
    | t => pure t
  pure (Ntrl (Var var j))
substNtrl (App t u) f = do
  t <- substNtrl t f
  u <- subst u f
  doApp t u
substNtrl Absurd f = pure $ Ntrl Absurd

subst (Ntrl t)  f = substNtrl t f
subst (Cnstr t) f = map Cnstr $ substCnstr t f
subst Irrel     f = pure Irrel

doApp (Ntrl t) u = pure $ Ntrl (App t u)
doApp Irrel u = pure $ Irrel
doApp (Cnstr (Lambda var t)) u = subst t (add (Logging ann . NormalForm) (pure u) (pure . point))
doApp (Cnstr t) u = inScope "wrong constructor in apply:" $ fatal t

partial
subst1 : NormalForm n -> NormalForm (S n) -> Logging ann (NormalForm n)
subst1 t u = subst u (add (Logging ann . NormalForm) (pure t) (pure . point))

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

-- invariant: all definitions fully expanded
-- invariant: m |- t, u <- a : s
covering partial
convert : (t, u, a : NormalForm n) -> Sort -> Logging ann Bool
-- In sort Prop
convert Irrel Irrel a Prop = pure True
-- In unknown type in set
convert t u (Ntrl _) (Set k) = pure $ t == u
-- In type Set
convert (Cnstr (Sort s')) (Cnstr (Sort s'')) (Cnstr (Sort _)) (Set _) = pure $ s' == s''
convert (Cnstr (Pi s s' _ a b)) (Cnstr (Pi l l' _ a' b')) (Cnstr (Sort _)) (Set _) =
  pure $
  s == l && s' == l' && !(convert a a' (cast s) (suc s)) && !(convert b b' (cast s') (suc s'))
convert (Cnstr Top) (Cnstr Top) (Cnstr (Sort _)) (Set _) = pure True
convert (Cnstr Bottom) (Cnstr Bottom) (Cnstr (Sort _)) (Set _) = pure True
convert (Ntrl t) u (Cnstr (Sort s)) (Set k) = pure $ Ntrl t == u
convert t (Ntrl u) (Cnstr (Sort s)) (Set k) = pure $ t == Ntrl u
convert t u (Cnstr (Sort s)) (Set k) = pure $ False
-- In type Pi
convert t u (Cnstr (Pi s s' var a b)) (Set k) = do
  t <- doApp (weaken 1 t) (Ntrl $ Var var FZ)
  u <- doApp (weaken 1 u) (Ntrl $ Var var FZ)
  convert t u b s'
-- Default
convert t u a s =
  inScope "invalid conversion:" $ fatal $
  fillSep {ann} [prettyPrec App t, prettyPrec App u, prettyPrec App a, prettyPrec App s]


-- 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, 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

covering partial
asSubst : (ctx : TyContext m n) -> Fin n -> Logging ann (NormalForm m)
asSubst (ctx :< def)           FZ     = subst def.tm (asSubst ctx)
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 ------------------------------------------------------

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

check ctx tm@(Lambda _ _ t) (Cnstr (Pi s s' var a b)) _ = do
  inScope "check" $ trace $ map (\_ => pretty {ann} "checking under lambda with type" <++> pretty a) (fullBounds tm)
  inScope "check" $ trace $ map (\_ => "binding new variable to \{var}") (fullBounds tm)
  inScope "check" $ trace $ map (\_ => pretty {ann} "checking for type" <++> pretty b) (fullBounds tm)
  t <- check (ctx ::< (var, a, s)) t b s'
  pure (Cnstr $ Lambda var t)
check ctx tm@(Lambda _ _ _) ty _ = typeMismatch (fullBounds tm).bounds (pretty "pi") (pretty ty)
check ctx tm ty s = do
  let bounded = fullBounds tm
  inScope "check" $ trace $ map (\_ => "checking has fallen through") bounded
  (tm, ty', s') <- infer ctx tm
  inScope "check" $ trace $ map (\_ => pretty {ann} "inferred type is" <++> pretty ty') bounded
  let True = s == s'
    | False => sortMismatch bounded.bounds (pretty s) (pretty s')
  True <- convert !(subst ty $ asSubst ctx) !(subst ty' $ asSubst ctx) (cast s) (suc s)
    | False => typeMismatch bounded.bounds (pretty ty) (pretty ty')
  inScope "check" $ trace $ map (\_ => pretty {ann} "converted" <++> pretty ty' <+> softline <+> pretty "to" <++> pretty ty) bounded
  pure tm

infer ctx tm@(Var b i) = do
  inScope "infer" $ trace $ map (\_ => "encountered variable \{b.val}@\{show i}") (fullBounds tm)
  let (t, a, s) = index ctx i
  inScope "infer" $ trace $ map (\_ => pretty {ann} "variable has type" <++> pretty a) (fullBounds tm)
  pure $ (t, a, s)
infer ctx (Sort b s) = pure $ (cast s, cast (suc s), suc (suc s))
infer ctx tm@(Pi _ var a b) = do
  inScope "infer" $ trace $ map (\_ => "encountered Pi type") (fullBounds tm)
  (a, s) <- inferType ctx a
  inScope "infer" $ trace $ map (\_ => pretty {ann} "argument has type" <++> pretty a) (fullBounds tm)
  (b, s') <- inferType (ctx ::< (var, 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')) (fullBounds tm)
  pure (Cnstr (Pi s s' var a b), cast (s ~> s'), suc (s ~> s'))
infer ctx tm@(Lambda _ _ _) = inScope "cannot infer type" $ fatal (fullBounds tm)
infer ctx tm@(App t u) = do
  inScope "infer" $ trace $ map (\_ => "encountered application") (fullBounds tm)
  let bounded = fullBounds t
  (t, ty@(Cnstr (Pi s s' _ a b)), _) <- infer ctx t
    | (_, t, _) => inScope "wrong type to apply" $ fatal (map (\_ => t) bounded)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "function has type" <++> pretty ty) (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "checking argument has type" <++> pretty a) (fullBounds tm)
  u <- check ctx u a s
  inScope "infer" $ trace $ map (\_ => "argument is well typed") (fullBounds 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) (fullBounds tm)
  pure (res, ty, s')
infer ctx (Top b) = pure $ (Cnstr Top, cast Prop, Set 0)
infer ctx (Point b) = pure $ (Irrel, Cnstr Top, Prop)
infer ctx (Bottom b) = pure $ (Cnstr Bottom, cast Prop, Set 0)
infer ctx tm@(Absurd b a t) = do
  inScope "infer" $ trace $ map (\_ => "encountered absurd") (fullBounds tm)
  (a, s) <- inferType ctx a
  inScope "infer" $ trace $ map (\_ => pretty {ann} "will fudge to type" <++> pretty a) (fullBounds tm)
  inScope "infer" $ trace $ map (\_ => pretty {ann} "checking for proof of false") (fullBounds tm)
  _ <- check ctx t (Cnstr Bottom) Prop
  inScope "infer" $ trace $ map (\_ => "proof of false is well typed") (fullBounds tm)
  pure (Ntrl Absurd, a, s)

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

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

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

covering 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)