{-# LANGUAGE TupleSections, OverloadedStrings #-}
module Elab where

import Elab.Monad

import qualified Presyntax.Presyntax as P

import Syntax
import Elab.Eval

infer :: P.Expr -> ElabM (Term, NFType)
infer (P.Var t) = (Ref (Bound t),) <$> getNfType (Bound t)

infer (P.App p f x) = do
  (f, f_ty) <- infer f
  (d, r, w) <- isPiType p f_ty
  x <- check x d
  x_nf <- eval x
  pure (App p (w f) x, r x_nf)

infer (P.Pi p s d r) = do
  d <- check d VType
  d_nf <- eval d
  assume (Bound s) d_nf $ do
    r <- check r VType
    pure (Pi p s d r, VType)

infer (P.Sigma s d r) = do
  d <- check d VType
  d_nf <- eval d
  assume (Bound s) d_nf $ do
    r <- check r VType
    pure (Sigma s d r, VType)

infer exp = do
  t <- newMeta VType
  tm <- check exp t
  pure (tm, t)

check :: P.Expr -> NFType -> ElabM Term
check (P.Lam p var body) (VPi p' dom (Closure _ rng)) | p == p' =
  assume (Bound var) dom $
    Lam p var <$> check body (rng (VVar (Bound var)))

check tm (VPi P.Im dom (Closure var rng)) =
  assume (Bound var) dom $
    Lam P.Im var <$> check tm (rng (VVar (Bound var)))

check (P.Lam p v b) ty = do
  (d, r, wp) <- isPiType p ty
  assume (Bound v) d $
    wp . Lam P.Im v <$> check b (r (VVar (Bound v)))

check (P.Pair a b) ty = do
  (d, r, wp) <- isSigmaType ty
  a <- check a d
  a_nf <- eval a
  b <- check b (r a_nf)
  pure (wp (Pair a b))

check exp ty = do
  (tm, has) <- infer exp
  unify has ty
  pure tm

isPiType :: P.Plicity -> NFType -> ElabM (Value, NFType -> NFType, Term -> Term)
isPiType p (VPi p' d (Closure _ k)) | p == p' = pure (d, k, id)
isPiType p t = do
  dom <- newMeta VType
  name <- newName
  assume (Bound name) dom $ do
    rng <- newMeta VType
    wp <- isConvertibleTo t (VPi p dom (Closure name (const rng)))
    pure (dom, const rng, wp)

isSigmaType :: NFType -> ElabM (Value, NFType -> NFType, Term -> Term)
isSigmaType (VSigma d (Closure _ k)) = pure (d, k, id)
isSigmaType t = do
  dom <- newMeta VType
  name <- newName
  assume (Bound name) dom $ do
    rng <- newMeta VType
    wp <- isConvertibleTo t (VSigma dom (Closure name (const rng)))
    pure (dom, const rng, wp)

identityTy :: NFType
identityTy = VPi P.Im VType (Closure "A" $ \t -> VPi P.Ex t (Closure "_" (const t)))