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

import Control.Arrow (Arrow(first))
import Control.Monad.Reader
import Control.Exception

import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Traversable
import Data.Text (Text)
import Data.Map (Map)
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

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 (P.LamCase pats) ty = do
  porp <- isPiType P.Ex ty
  case porp of
    It'sProd dom rng wp -> do
      name <- newName
      cases <- for pats $ \(pat, rhs) -> do
        checkPattern pat dom \pat wp pat_nf -> do
          rhs <- check rhs (rng pat_nf)
          pure (pat, wp rhs)
      let x = wp (Lam P.Ex name (Case (Ref name) cases))
      pure x
    _ -> do
      dom <- newMeta VTypeω
      n <- newName' (Bound (T.singleton 'x') 0)
      assume n dom \_ -> do
        rng <- newMeta VTypeω
        throwElab $ NotEqual (VPi P.Ex dom (Closure n (const rng))) ty

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

checkPattern :: P.Pattern -> NFSort -> (Term -> (Term -> Term) -> Value -> ElabM a) -> ElabM a
checkPattern (P.PCap var) dom cont = do
  name <- asks (Map.lookup var . nameMap)
  case name of
    Just name@(ConName _ _ skip arity) -> do
      unless (arity == 0) $ throwElab $ UnsaturatedCon name
      (ty, wp) <- instantiate =<< getNfType name
      unify ty dom
      wrap <- skipLams skip
      cont (Con name) wrap =<< eval (wp (Con name))
    Just name -> throwElab $ NotACon name
    Nothing -> assume (Bound var 0) dom \name -> cont (Ref name) (Lam P.Ex name) (VVar name)
checkPattern (P.PCon var args) dom cont =
  do
    name <- asks (Map.lookup var . nameMap)
    case name of
      Just name@(ConName _ _ nskip arity) -> do
        unless (arity == length args) $ throwElab $ UnsaturatedCon name
        (ty, wp) <- instantiate =<< getNfType name
        _ <- isConvertibleTo (skipBinders arity ty) dom
        skip <- skipLams nskip
        bindNames args ty $ \names wrap ->
          cont (Con name) (skip . wrap) =<< eval (foldl (\x y -> App P.Ex x (Ref y)) (wp (Con name)) names)
      Just name -> throwElab $ NotACon name
      _ -> throwElab $ NotInScope (Bound var 0)
  where
    skipBinders :: Int -> NFType -> NFType
    skipBinders 0 t = t
    skipBinders n (VPi _ _ (Closure v r)) = skipBinders (n - 1) (r (VVar v))
    skipBinders _ _ = error $ "constructor type is wrong?"

    bindNames (n:ns) (VPi p d (Closure _ r)) k =
      assume (Bound n 0) d \n -> bindNames ns (r (VVar n)) \ns w ->
        k (n:ns) (Lam p n . w)
    bindNames [] _ k = k [] id
    bindNames xs t _ = error $ show (xs, t)

instantiate :: NFType -> ElabM (NFType, Term -> Term)
instantiate (VPi P.Im d (Closure _ k)) = do
  t <- newMeta d
  (ty, w) <- instantiate (k t)
  pure (ty, \inner -> App P.Im (w inner) (quote t))
instantiate x = pure (x, id)

skipLams :: Int -> ElabM (Term -> Term)
skipLams 0 = pure id
skipLams k = do
  n <- newName
  (Lam P.Im n . ) <$> skipLams (k - 1)

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
      makeLetDef (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
      makeLetDef (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
      makeLetDef (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 <- evalFix (Defined name 0) ty_nf rhs
      makeLetDef (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

checkStatement (P.Data name tele retk constrs) k =
  do
    checkTeleRetk True tele retk \kind tele undef -> do
      kind_nf <- eval kind
      defineInternal (Defined name 0) kind_nf (\name' -> VNe (HData name') mempty) \name' -> 
        checkCons tele (VNe (HData name') (Seq.fromList (map makeProj tele))) constrs (local (markAsDef name' . undef) k)
  where
    makeProj (x, p, _) = PApp p (VVar x)

    markAsDef x e = e { definedNames = Set.insert x (definedNames e) }

    checkTeleRetk allKan [] retk cont = do
      t <- check retk VTypeω
      t_nf <- eval t
      when allKan $ unify t_nf VType
      cont t [] id
    checkTeleRetk allKan ((x, p, t):xs) retk cont = do
      (t, ty) <- infer t
      _ <- isConvertibleTo ty VTypeω
      let
        allKan' = case ty of
          VType -> allKan
          _ -> False
      t_nf <- eval t
      let rm nm e = e{ nameMap = Map.delete (getNameText nm) (nameMap e), getEnv = Map.delete nm (getEnv e) }
      assume (Bound x 0) t_nf $ \nm -> checkTeleRetk allKan' xs retk \k xs w -> cont (Pi p nm t k) ((nm, p, t_nf):xs) (rm nm . w)

    checkCons _ _et [] k = k

    checkCons n ret ((x, ty):xs) k = do
      t <- check ty VTypeω
      ty_nf <- eval t
      let
        (args, ret') = splitPi ty_nf
        closed = close n t
        n' = map (\(x, _, y) -> (x, P.Im, y)) n
      unify ret' ret
      closed_nf <- eval closed
      defineInternal (ConName x 0 (length n') (length args)) closed_nf (makeCon closed_nf mempty n' args) \_ -> checkCons n ret xs k

    close []             t = t
    close ((x, _, y):xs) t = Pi P.Im x (quote y) (close xs t)

    splitPi (VPi p y (Closure x k)) = first ((x, p, y):) $ splitPi (k (VVar x))
    splitPi t = ([], t)

    makeCon cty sp []              [] con = VNe (HCon cty con) sp
    makeCon cty sp ((nm, p, _):xs) ys con = VLam p $ Closure nm \a -> makeCon cty (sp Seq.:|> PApp p a) xs ys con
    makeCon cty sp [] ((nm, p, _):ys) con = VLam p $ Closure nm \a -> makeCon cty (sp Seq.:|> PApp p a) [] ys con

evalFix :: Name -> NFType -> Term -> ElabM Value
evalFix name nft term = do
  env <- ask
  pure . fix $ \val -> eval' env{ getEnv = Map.insert name (nft, val) (getEnv env) } term

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)

data UnsaturatedCon = UnsaturatedCon { theConstr :: Name }
  deriving          (Show, Typeable)
  deriving anyclass (Exception)

data NotACon = NotACon { theNotConstr :: Name }
  deriving          (Show, Typeable)
  deriving anyclass (Exception)