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

-- Errors ----------------------------------------------------------------------

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

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

report : Bounds -> Doc ann -> Error ann a
report = curry Left . Just

internal : Doc ann -> Error ann a
internal = Left . (Nothing, )

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

PointedRename (Error ann . NormalForm) where
  point = pure . point

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

wkn : (k : _) -> (Fin n -> Error ann (NormalForm m)) -> Fin (k + n) -> Error ann (NormalForm (k + m))
wkn = lift {x = Error ann . NormalForm}

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

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

substNtrl (Var i) f = f i
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 t)) u = subst t (add (Error ann . NormalForm) (pure u) (pure . point))
doApp (Cnstr t) u = internal (pretty "wrong constructor in apply:" <++> pretty t)

partial
subst1 : NormalForm n -> NormalForm (S n) -> Error ann (NormalForm n)
subst1 t u = subst u (add (Error 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 -> Error 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' a b)) (Set k) = do
  t <- doApp (weaken 1 t) (point FZ)
  u <- doApp (weaken 1 u) (point FZ)
  convert t u b s'
-- Default
convert t u a s = internal $ fillSep ["invalid convert:", 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 ::< (_, ty, s)) FZ     = (point 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 -> Error ann (NormalForm m)
asSubst (ctx :< def) FZ     = subst def.tm (asSubst ctx)
asSubst (ctx ::< _)  FZ     = pure $ point 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 -> Error ann (NormalForm n)
covering partial
infer : TyContext m n -> Term n -> Error ann (NormalForm n, NormalForm n, Sort)
covering partial
inferSort : {default "type mismatch" tag : String}
  -> TyContext m n -> Term n -> Error ann (NormalForm n, Sort)

check ctx (Lambda _ var t) (Cnstr (Pi s s' a b)) _ = do
  t <- check (ctx ::< (var, a, s)) t b s'
  pure (Cnstr $ Lambda t)
check ctx tm@(Lambda _ _ _) ty _ =
  report (fullBounds tm).bounds
    (pretty "type mismatch: expected pi" <+> comma <+> softline <+> "got" <++> pretty ty)
check ctx tm ty s = do
  let bounds = (fullBounds tm).bounds
  (tm, ty', s') <- infer ctx tm
  let True = s == s'
    | False => report bounds
      (pretty "sort mismatch: expected" <++> pretty s <+> comma <+> softline <+> "got" <++> pretty s')
  True <- convert !(subst ty $ asSubst ctx) !(subst ty' $ asSubst ctx) (cast 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 $ (cast s, cast (suc s), suc (suc s))
infer ctx (Pi _ var a b) = do
  (a, s) <- inferSort ctx a
  let aDef = MkDefinition {name = var, sort = suc s, ty = cast s, tm = a}
  (b, s') <- inferSort (ctx :< aDef) b
  pure (Cnstr (Pi s s' a b), cast (s ~> s'), suc (s ~> s'))
infer ctx tm@(Lambda _ _ _) = report (fullBounds tm).bounds (pretty "cannot infer type:" <+> softline <+> pretty tm)
infer ctx (App t u) = do
  let bounds = (fullBounds t).bounds
  (t, Cnstr (Pi s s' a b), _) <- infer ctx t
    | (_, t, _) => report bounds (pretty "wrong type to apply:" <+> softline <+> pretty t)
  u <- check ctx u a s
  res <- doApp t u
  ty <- subst1 t b
  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 (Absurd b a t) = do
  (a, s) <- inferSort ctx a
  _ <- check ctx t (Cnstr Bottom) Prop
  pure (Ntrl Absurd, a, s)

inferSort ctx a = do
  (a, Cnstr (Sort s), _) <- infer ctx a
    | (_, ty, _) => report (fullBounds a).bounds
      (pretty "\{tag}: expected sort" <+> comma <+> softline <+> "got" <++> pretty ty)
  pure (a, s)

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

covering partial
checkDef : Context n -> Term.Definition n -> Error ann (NormalForm.Definition n)
checkDef ctx def = do
  let tyBounds = (fullBounds def.ty).bounds
  let tmBounds = (fullBounds def.tm).bounds
  (ty, sort) <- inferSort {tag = "invalid declaration"} (fromContext ctx) def.ty
  tm <- check (fromContext ctx) def.tm ty sort
  pure $ MkDefinition {name = def.name, ty, tm, sort}

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