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

import Control.Monad.Reader
import Control.Exception

import qualified Data.Map.Strict as Map
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
import Syntax.Pretty
import qualified Data.Text as T

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

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

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

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 (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) (VIsOne phi) $ do
            env <- ask
            for (truthAssignments phi (getEnv env)) $ \e -> do
              let env' = env{ getEnv = e }
              check rhs (eval' env' dom_q)
          Nothing -> do
            env <- ask
            for (truthAssignments phi (getEnv env)) $ \e -> do
              let env' = env{ getEnv = e }
              check rhs (eval' env' dom_q)
    pure (n, (phi, (P.condV formula, 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 (Bound v) (VIsOne formula)
        Nothing -> id
    k $ for_ eqns $ \(n', (formula', (binding, rhs'))) -> do
      let
        k = case binding of
          Just v -> assume (Bound v) (VIsOne formula)
          Nothing -> id
        truth = possible mempty (iand formula formula')
        add [] = id
        add ((~(HVar x), True):xs) = define x VI VI1 . add xs
        add ((~(HVar x), False):xs) = define 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 (Bound name) y)) eqns))))

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

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 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 (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 (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)

isPartialType :: NFType -> ElabM (NFEndp, Value)
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 (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 ()
        _ -> throwElab $ 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
      _ -> 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) (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)