{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Elab where

import Control.Monad.Reader
import Control.Exception

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Data.Traversable
import Data.Typeable
import Data.Foldable

import Elab.Eval.Formula (possible, truthAssignments)
import Elab.WiredIn
import Elab.Monad
import Elab.Eval

import qualified Presyntax.Presyntax as P

import Prettyprinter

import Syntax.Pretty
import Syntax
import Data.Map (Map)
import Data.Text (Text)

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)
    It'sPartial phi a w -> do
      x <- check x (VIsOne phi)
      pure (App P.Ex (w f) x, a)
    It'sPartialP phi a w -> do
      x <- check x (VIsOne phi)
      x_nf <- eval x
      pure (App P.Ex (w f) x, a @@ 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 0) dom $ \name ->
    Lam p name <$> check body (rng (VVar name))

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

check (P.Lam p v b) ty = do
  porp <- isPiType p =<< forceIO ty
  case porp of
    It'sProd d r wp ->
      assume (Bound v 0) d $ \name ->
        wp . Lam p name <$> check b (r (VVar name))

    It'sPath li le ri wp -> do
      tm <- assume (Bound v 0) VI $ \var ->
        Lam P.Ex var <$> check b (force (li (VVar var)))

      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)) (quote le) (quote ri) tm))

    It'sPartial phi a wp ->
      assume (Bound v 0) (VIsOne phi) $ \var ->
        wp . Lam p var <$> check b a

    It'sPartialP phi a wp ->
      assume (Bound v 0) (VIsOne phi) $ \var ->
        wp . Lam p var <$> check b (a @@ VVar var)

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 0) d_nf \var -> do
    r <- check r ty
    pure (Pi p var d r)

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

check (P.Let items body) ty = do
  checkLetItems mempty items \decs -> do
    body <- check body ty
    pure (Let decs body)

check (P.LamSystem bs) ty = do
  (extent, dom) <- isPartialType ty
  let dom_q = quote dom
  eqns <- for (zip [(0 :: Int)..] bs) $ \(n, (formula, rhs)) -> do
    phi <- checkFormula (P.condF formula)
    rhses <-
        case P.condV formula of
          Just t -> assume (Bound t 0) (VIsOne phi) $ \var -> do
            env <- ask
            for (truthAssignments phi (getEnv env)) $ \e -> do
              let env' = env{ getEnv = e }
              (Just var,) <$> check rhs (eval' env' dom_q)
          Nothing -> do
            env <- ask
            for (truthAssignments phi (getEnv env)) $ \e -> do
              let env' = env{ getEnv = e }
              (Nothing,) <$> check rhs (eval' env' dom_q)
    pure (n, (phi, head rhses))

  unify extent (foldl ior VI0 (map (fst . snd) eqns))

  for_ eqns $ \(n, (formula, (binding, rhs))) -> do
    let
      k = case binding of
        Just v -> assume v (VIsOne formula) . const
        Nothing -> id
    k $ for_ eqns $ \(n', (formula', (binding, rhs'))) -> do
      let
        k = case binding of
          Just v -> assume v (VIsOne formula) . const
          Nothing -> id
        truth = possible mempty (iand formula formula')
        add [] = id
        add ((~(HVar x), True):xs) = redefine x VI VI1 . add xs
        add ((~(HVar x), False):xs) = redefine x VI VI0 . add xs
      k $ when ((n /= n') && fst truth) . add (Map.toList (snd truth)) $ do
        vl <- eval rhs
        vl' <- eval rhs'
        unify vl vl'
          `withNote` vsep [ pretty "These two cases must agree because they are both possible:"
                          , indent 2 $ pretty '*' <+> prettyTm (quote formula) <+> operator (pretty "=>") <+> prettyTm rhs
                          , indent 2 $ pretty '*' <+> prettyTm (quote formula') <+> operator (pretty "=>") <+> prettyTm rhs'
                          ]
          `withNote` (pretty "Consider this face, where both are true:" <+> showFace (snd truth))

  name <- newName
  let
    mkB name (Just v, b) = App P.Ex (Lam P.Ex v b) (Ref name)
    mkB _ (Nothing, b) = b
  pure (Lam P.Ex name (System (Map.fromList (map (\(_, (x, y)) -> (quote x, mkB name y)) eqns))))

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

checkLetItems :: Map Text (Maybe NFType) -> [P.LetItem] -> ([(Name, Term, Term)] -> ElabM a) -> ElabM a
checkLetItems _   [] cont = cont []
checkLetItems map (P.LetDecl v t:xs) cont = do
  t <- check t VTypeω
  t_nf <- eval t
  assume (Defined v 0) t_nf \_ ->
    checkLetItems (Map.insert v (Just t_nf) map) xs cont

checkLetItems map (P.LetBind name rhs:xs) cont = do
  case Map.lookup name map of
    Nothing -> do
      (tm, ty) <- infer rhs
      tm_nf <- eval tm
      define (Defined name 0) ty tm_nf \name' ->
        checkLetItems map xs \xs ->
          cont ((name', quote ty, tm):xs)
    Just Nothing -> throwElab $ Redefinition (Defined name 0)
    Just (Just ty_nf) -> do
      rhs <- check rhs ty_nf
      rhs_nf <- eval rhs
      define (Defined name 0) ty_nf rhs_nf \name' ->
        checkLetItems (Map.insert name Nothing map) xs \xs ->
          cont ((name', quote ty_nf, rhs):xs)

checkFormula :: P.Formula -> ElabM Value
checkFormula P.FTop = pure VI1
checkFormula P.FBot = pure VI0
checkFormula (P.FAnd x y) = iand <$> checkFormula x <*> checkFormula y
checkFormula (P.FOr x y) = ior <$> checkFormula x <*> checkFormula y
checkFormula (P.FIs0 x) = do
  nm <- getNameFor x
  ty <- getNfType nm
  unify ty VI
  pure (inot (VVar nm))
checkFormula (P.FIs1 x) = do
  nm <- getNameFor x
  ty <- getNfType nm
  unify ty VI
  pure (VVar nm)

isSort :: NFType -> ElabM ()
isSort t = isSort (force t) where
  isSort VType              = pure ()
  isSort VTypeω             = pure ()
  isSort vt@(VNe HMeta{} _) = unify vt VType
  isSort vt = throwElab $ 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
             }
  | It'sPartial { partialExtent :: NFEndp
                , partialDomain :: Value
                , partialWrap   :: Term -> Term
                }
  | It'sPartialP { partialExtent :: NFEndp
                 , partialDomain :: Value
                 , partialWrap   :: Term -> Term
                 }

isPiType :: P.Plicity -> NFType -> ElabM ProdOrPath
isPiType p x = isPiType p (force x) where
  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 (VPartial phi a) = pure (It'sPartial phi a id)
  isPiType P.Ex (VPartialP phi a) = pure (It'sPartialP phi a 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)))
      It'sPartial phi a w -> It'sPartial phi a (\f -> w (App P.Im f (quote meta)))
      It'sPartialP phi a w -> It'sPartialP phi a (\f -> w (App P.Im f (quote meta)))
  isPiType p t = do
    dom <- newMeta VType
    name <- newName
    assume name dom $ \name -> 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 t = isSigmaType (force t) where
  isSigmaType (VSigma d (Closure _ k)) = pure (d, k, id)
  isSigmaType t = do
    dom <- newMeta VType
    name <- newName
    assume name dom $ \name -> do
      rng <- newMeta VType
      wp <- isConvertibleTo t (VSigma dom (Closure name (const rng)))
      pure (dom, const rng, wp)

isPartialType :: NFType -> ElabM (NFEndp, Value)
isPartialType t = isPartialType (force t) where
  isPartialType (VPartial phi a) = pure (phi, a)
  isPartialType (VPartialP phi a) = pure (phi, a)
  isPartialType t = do
    phi <- newMeta VI
    dom <- newMeta (VPartial phi VType)
    unify t (VPartial phi dom)
    pure (phi, dom)

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 (flip Defined 0 <$> name) ty_nf (const k)

checkStatement (P.Postulate []) k = k
checkStatement (P.Postulate ((name, ty):xs)) k = do
  ty <- check ty VTypeω
  ty_nf <- eval ty
  assume (Defined name 0) ty_nf \name ->
    local (\e -> e { definedNames = Set.insert name (definedNames e) }) (checkStatement (P.Postulate xs) k)

checkStatement (P.Defn name rhs) k = do
  ty <- asks (Map.lookup name . nameMap)
  case ty of
    Nothing -> do
      (tm, ty) <- infer rhs
      tm_nf <- eval tm
      define (Defined name 0) ty tm_nf (const k)
    Just nm -> do
      ty_nf <- getNfType nm
      t <- asks (Set.member nm . definedNames)
      when t $ throwElab (Redefinition (Defined name 0))

      rhs <- check rhs ty_nf
      rhs_nf <- eval rhs
      define (Defined name 0) ty_nf rhs_nf $ \name ->
        local (\e -> e { definedNames = Set.insert name (definedNames e) }) k

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

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

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

  define (Defined var 0) (wiType wi) (wiValue wi) $ \name ->
    local (\e -> e { definedNames = Set.insert name (definedNames e) }) 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)