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

import Control.Monad.Reader
import Control.Exception

import qualified Data.Map.Strict as Map
import Data.Typeable

import Elab.WiredIn
import Elab.Monad
import Elab.Eval

import qualified Presyntax.Presyntax as P

import Prettyprinter

import Syntax.Pretty
import Syntax

infer :: P.Expr -> ElabM (Term, NFType)
infer (P.Span ex a b) = withSpan a b $ infer ex

infer (P.Var t) = do
  name <- getNameFor t
  nft <- getNfType name
  pure (Ref name, nft)

infer (P.App p f x) = do
  (f, f_ty) <- infer f
  porp <- isPiType p f_ty
  case porp of
    It'sProd d r w -> do
      x <- check x d
      x_nf <- eval x
      pure (App p (w f) x, r x_nf)
    It'sPath li le ri wp -> do
      x <- check x VI
      x_nf <- eval x
      pure (IElim (quote (fun li)) (quote le) (quote ri) (wp f) x, li x_nf)

infer (P.Proj1 x) = do
  (tm, ty) <- infer x
  (d, _, wp) <- isSigmaType ty
  pure (Proj1 (wp tm), d)

infer (P.Proj2 x) = do
  (tm, ty) <- infer x
  tm_nf <- eval tm
  (_, r, wp) <- isSigmaType ty
  pure (Proj2 (wp tm), r (vProj1 tm_nf))

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

check :: P.Expr -> NFType -> ElabM Term
check (P.Span ex a b) ty = withSpan a b $ check ex ty

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
  porp <- isPiType p =<< forceIO ty
  case porp of
    It'sProd d r wp -> do
      assume (Bound v) d $
        wp . Lam p v <$> check b (r (VVar (Bound v)))
    It'sPath li le ri wp -> do
      tm <- assume (Bound v) VI $
        Lam P.Ex v <$> check b (li (VVar (Bound v)))
      tm_nf <- eval tm
      unify (tm_nf @@ VI0) le
        `catchElab` (throwElab . WhenCheckingEndpoint le ri VI0)
      unify (tm_nf @@ VI1) ri
        `catchElab` (throwElab . WhenCheckingEndpoint le ri VI1)
      pure (wp (PathIntro (quote (fun li)) tm))

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

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

check (P.Sigma s d r) ty = do
  isSort ty
  d <- check d ty
  d_nf <- eval d
  assume (Bound s) d_nf $ do
    r <- check r ty
    pure (Sigma s d r)

check exp ty = do
  (tm, has) <- switch $ infer exp
  wp <- isConvertibleTo has ty
  pure (wp tm)

isSort :: NFType -> ElabM ()
isSort VType                = pure ()
isSort VTypeω               = pure ()
isSort vt@(VNe HMeta{} _)   = unify vt VType
isSort vt = liftIO . throwIO $ NotEqual vt VType

data ProdOrPath
  = It'sProd { prodDmn  :: NFType
             , prodRng  :: NFType -> NFType
             , prodWrap :: Term -> Term
             }
  | It'sPath { pathLine  :: NFType -> NFType
             , pathLeft  :: Value
             , pathRight :: Value
             , pathWrap  :: Term -> Term
             }

isPiType :: P.Plicity -> NFType -> ElabM ProdOrPath
isPiType p (VPi p' d (Closure _ k)) | p == p' = pure (It'sProd d k id)
isPiType P.Ex (VPath li le ri) = pure (It'sPath (li @@) le ri id)
isPiType P.Ex (VPi P.Im d (Closure _ k)) = do
  meta <- newMeta d
  porp <- isPiType P.Ex (k meta)
  pure $ case porp of
    It'sProd d r w -> It'sProd d r (\f -> w (App P.Im f (quote meta)))
    It'sPath l x y w -> It'sPath l x y (\f -> w (App P.Im f (quote meta)))
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 (It'sProd 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)))

checkStatement :: P.Statement -> ElabM a -> ElabM a
checkStatement (P.SpanSt s a b) k = withSpan a b $ checkStatement s k

checkStatement (P.Decl name ty) k = do
  ty <- check ty VTypeω
  ty_nf <- eval ty
  assumes (Defined <$> name) ty_nf k

checkStatement (P.Defn name rhs) k = do
  ty <- asks (Map.lookup (Defined name) . getEnv)
  case ty of
    Nothing -> do
      (tm, ty) <- infer rhs
      tm_nf <- eval tm
      define (Defined name) ty tm_nf k
    Just (ty_nf, nm) -> do
      case nm of
        VVar (Defined n') | n' == name -> pure ()
        _ -> liftIO . throwIO $ Redefinition (Defined name)

      rhs <- check rhs ty_nf
      rhs_nf <- eval rhs
      define (Defined name) ty_nf rhs_nf k

checkStatement (P.Builtin winame var) k = do
  wi <-
    case Map.lookup winame wiredInNames of
      Just wi -> pure wi
      _ -> liftIO . throwIO $ NoSuchPrimitive winame

  let
    check = do
      nm <- getNameFor var
      ty <- getNfType nm
      unify ty (wiType wi)
        `withNote` hsep [ "Previous definition of", pretty nm, "here" ]
        `seeAlso` nm

  env <- ask
  liftIO $
    runElab check env `catch` \(_ :: NotInScope) -> pure ()

  define (Defined var) (wiType wi) (wiValue wi) k

checkStatement (P.ReplNf e) k = do
  (e, _) <- infer e
  e_nf <- eval e
  h <- asks commHook
  liftIO (h e_nf)
  k

checkStatement (P.ReplTy e) k = do
  (_, ty) <- infer e
  h <- asks commHook
  liftIO (h ty)
  k

checkProgram :: [P.Statement] -> ElabM a -> ElabM a
checkProgram [] k = k
checkProgram (st:sts) k = checkStatement st $ checkProgram sts k

newtype Redefinition = Redefinition { getRedefName :: Name }
  deriving (Show, Typeable, Exception)

data WhenCheckingEndpoint = WhenCheckingEndpoint { leftEndp :: Value, rightEndp :: Value, whichIsWrong :: NFEndp, exc :: SomeException }
  deriving (Show, Typeable, Exception)