module CC.Term.Elaborate

import CC.Name
import CC.Term
import CC.Term.Eval
import CC.Term.Raw
import CC.Thinning

import Data.DPair
import Data.List1
import Data.Singleton
import Data.SnocList
import Data.SnocList.Elem

import Text.Bounded

%prefix_record_projections off

-- Elaboration Errors ----------------------------------------------------------

export
data ElabError : Type -> Type where
  Ok : ty -> ElabError ty
  Error : List1 (Either Name (String, Maybe Bounds)) -> ElabError ty

Functor ElabError where
  map f (Ok x) = Ok (f x)
  map f (Error xs) = Error xs

Applicative ElabError where
  pure = Ok

  Ok f <*> Ok v = Ok (f v)
  Ok f <*> Error ys = Error ys
  Error xs <*> Ok v = Error xs
  Error xs <*> Error ys = Error (xs ++ ys)

Monad ElabError where
  join (Ok (Ok x)) = Ok x
  join (Ok (Error xs)) = Error xs
  join (Error xs) = Error xs

export
runElabError : ElabError ty -> Either (List1 (Either Name (String, Maybe Bounds))) ty
runElabError (Ok x) = Right x
runElabError (Error xs) = Left xs

fromNameError : NameError ty -> ElabError ty
fromNameError v =
  case runNameError v of
    Left errs => Error (map Left errs)
    Right v => Ok v

-- Elaboration Context ---------------------------------------------------------

record ElabCtx (local, global : Context) where
  constructor MkCtx
  env : Env local global
  types : (k : Nat) -> forall n. (0 prf : IsVar k n local) -> Lazy (Value global)
  pos : Maybe Bounds

%name ElabCtx ctx

wknCtx :
  {global : Context} ->
  ElabCtx local global ->
  (n : Name) ->
  ElabCtx local (global :< n)
wknCtx ctx n = { env $= flip wknEnv (drop id), types := types } ctx
  where
  types : (k : Nat) -> forall m. (0 _ : IsVar k m local) -> Lazy (Value (global :< n))
  types k prf = wknVal (ctx.types k prf) (drop id)

define :
  ElabCtx local global ->
  (n : Name) ->
  (val : Lazy (Value global)) ->
  (ty : Lazy (Value global)) ->
  ElabCtx (local :< n) global
define ctx n val ty = { env $= (:< (n, val)), types := types } ctx
  where
  types : (k : Nat) -> forall m. (0 _ : IsVar k m (local :< n)) -> Lazy (Value global)
  types 0 (Here eq) = ty
  types (S k) (There prf) = ctx.types k prf

bind :
  {global : Context} ->
  ElabCtx local global ->
  (n : Name) ->
  (ty : Lazy (Value global)) ->
  ElabCtx (local :< n) (global :< n)
bind ctx n ty = define (wknCtx ctx n) n (Ntrl $ VVar 0 _ $ fromElem Here) (wknVal ty (drop id))

updatePos : ElabCtx local global -> (irrel : Bool) -> Bounds -> ElabCtx local global
updatePos ctx irrel bounds = { pos := if irrel then Nothing else Just bounds } ctx

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

record PiType (ctx : Context) where
  constructor MkPi
  {0 local : Context}
  name : Name
  dom : Lazy (Value ctx)
  env : Env local ctx
  cod : Term (local :< name)

check :
  {local, global : Context} ->
  ElabCtx local global ->
  RawTerm ->
  Lazy (Value global) ->
  ElabError (Term local)
checkType :
  {local, global : Context} ->
  ElabCtx local global ->
  RawTerm ->
  ElabError (Term local)
infer :
  {local, global : Context} ->
  ElabCtx local global ->
  RawTerm ->
  ElabError (Term local, Lazy (Value global))
inferPi :
  {local, global : Context} ->
  ElabCtx local global ->
  RawTerm ->
  ElabError (Term local, PiType global)
inferType :
  {local, global : Context} ->
  ElabCtx local global ->
  RawTerm ->
  ElabError (Term local)

check ctx (Bounded (MkBounded t irrel bounds)) ty = check (updatePos ctx irrel bounds) t ty

check ctx (Let n Nothing val t) ty = do
  (val, ty') <- infer ctx val
  t <- check (define ctx n (eval ctx.env val) ty') t ty
  pure $ Let {n, ty = Nothing, val, t}
check ctx (Let n (Just ty') val t) ty = do
  ty' <- checkType ctx ty'
  let vty : Lazy (Value global) = eval ctx.env ty'
  val <- check ctx val vty
  t <- check (define ctx n (eval ctx.env val) vty) t ty
  pure $ Let {n, ty = Just ty', val, t}

check ctx Set VSet = pure $ Set

check ctx (Abs Nothing t) (VFunc env dom cod) = do
  t <- check ctx t (eval env cod)
  pure $ Abs (fresh local) (wknTerm t (drop id))
check ctx (Abs (Just n) t) (VFunc env dom cod) = do
  t <- check (bind ctx n (eval env dom)) t (eval (wknEnv env (drop id)) cod)
  pure $ Abs n t
check ctx (Abs Nothing t) (VPi m dom env cod) = do
  t <- check (wknCtx ctx m) t (eval (liftEnv [<m] env) cod)
  pure $ Abs m (wknTerm t (drop id))
check ctx (Abs (Just n) t) (VPi m dom env cod) = do
  t <- check (bind ctx n dom) t (eval (liftEnv [<n] env) $ renTerm Refl cod)
  pure $ Abs n t

check ctx t ty = do
  (t, ty') <- infer ctx t
  let True = convert ty ty'
    | False => Error $ singleton $ Right ("conversion failed", ctx.pos)
  pure t

checkType ctx t = check ctx t VSet

infer ctx (Bounded (MkBounded t irrel bounds)) = infer (updatePos ctx irrel bounds) t
infer ctx (Var n) = do
  Element k prf <- fromNameError $ lookup local n
  pure (Var k n prf, ctx.types k prf)
infer ctx (Let n Nothing val t) = do
  (val, ty') <- infer ctx val
  (t, ty) <- infer (define ctx n (eval ctx.env val) ty') t
  pure $ (Let {n, ty = Nothing, val, t}, ty)
infer ctx (Let n (Just ty) val t) = do
  ty' <- checkType ctx ty
  let vty : Lazy (Value global) = eval ctx.env ty'
  val <- check ctx val vty
  (t, ty) <- infer (define ctx n (eval ctx.env val) vty) t
  pure $ (Let {n, ty = Just ty', val, t}, ty)
infer ctx Set = pure (Set, delay VSet)
infer ctx (Pi (Nothing `Of` dom) cod) = do
  (dom, cod) <- [| MkPair (inferType ctx dom) (inferType ctx cod) |]
  pure $ (Pi Nothing dom cod, delay VSet)
infer ctx (Pi (Just n `Of` dom) cod) = do
  dom <- inferType ctx dom
  let vdom : Lazy (Value global) = eval ctx.env dom
  cod <- inferType (bind ctx n vdom) cod
  pure $ (Pi (Just n) dom cod, delay VSet)
infer ctx (App t u) = do
  (t, ty) <- inferPi ctx t
  u <- check ctx u ty.dom
  let vu : Lazy (Value global) = eval ctx.env u
  pure $ (App t u, delay $ eval (ty.env :< (ty.name, vu)) ty.cod)
infer ctx t = Error $ singleton $ Right ("cannot infer type", ctx.pos)

inferPi ctx t = do
  (t, ty) <- infer ctx t
  case force ty of
    VPi n dom env cod =>
      pure $ (t, MkPi n dom env cod)
    VFunc env dom cod =>
      let local' = recomputeLocal env in
      pure $ (t, MkPi (fresh' local') (eval env dom) env (wknTerm cod (drop $ id' local')))
    ty => Error $ singleton $ Right ("expected pi type", ctx.pos)

inferType ctx t = do
  (t, ty) <- infer ctx t
  case force ty of
    VSet => pure t
    ty => Error $ singleton $ Right ("expected set type", ctx.pos)

export
inferTerm : RawTerm -> ElabError (Term [<], Lazy (Value [<]))
inferTerm = infer (MkCtx [<] types Nothing)
  where
  types : (k : Nat) -> forall n. (0 _ : IsVar k n [<]) -> Lazy (Value [<])
  types k prf impossible