{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DerivingVia #-}
module Elaboration where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Concurrent

import qualified Data.HashMap.Strict as HashMap
import qualified Data.IntMap.Strict as IntMap
import Data.Text (Text)

import Elaboration.Monad

import Evaluate

import Presyntax

import Syntax

import System.IO.Unsafe ( unsafeDupablePerformIO )

import Value

elabNext :: MVar Int
elabNext = unsafeDupablePerformIO (newMVar 0)
{-# NOINLINE elabNext #-}

freshMeta :: Value -> ElabM Term
freshMeta expected = do
  ctx <- ask
  names <- getNames
  thisMeta <- liftIO $ do
    m <- modifyMVar elabNext $ \x -> pure (x + 1, x)
    modifyMVar_ elabMetas $ pure . IntMap.insert m (Unsolved names expected)
    pure m
  pure $ NewMeta (MV thisMeta) (elabBound ctx)

insert :: Term -> VTy -> ElabM (Term, VTy)
insert f (VPi Im _ d r) = do
  t <- freshMeta d
  t_nf <- asks (flip evaluate t . elabEnv)
  insert (App Im f t) (r $$ t_nf)
insert f x = pure (f, x)

insert' :: Term -> VTy -> ElabM (Term, VTy)
insert' t@(Lam Im _ _) ty = pure (t, ty)
insert' t ty = insert t ty

infer :: RawExpr -> ElabM (Term, VTy)
infer (RSrcPos start end expr) = local (\st -> st { elabSourcePos = (start, end) }) (infer expr)

infer (Rvar name) = ask >>= lookup where
  lookup ElabState{elabNames, elabConstrs, elabLevel} =
    case HashMap.lookup name elabNames of
      Just (l, t) -> pure (Bv (lvl2Ix elabLevel l), t)
      Nothing ->
        case HashMap.lookup name elabConstrs of
          Just t  -> pure (Con name, t)
          Nothing -> typeError (NotInScope name)

infer (Rapp p x y) = do
  (x, x_ty) <-
    infer x >>= \(x, x_ty) ->
      case p of
        Ex -> insert x x_ty
        _  -> pure (x, x_ty)

  (_, d, r) <- isPiType p x_ty
  y         <- check y d
  y_nf      <- asks (flip evaluate y . elabEnv)

  pure (App p x y, r $$ y_nf)

infer (Rpi e v d r) = do
  d <- check d VType
  d_nf <- asks (flip evaluate d . elabEnv)
  assumeLocal v d_nf $ do
    r <- check r VType
    pure (Pi e v d r, VType)

infer (Rsigma v d r) = do
  d <- check d VType
  d_nf <- asks (flip evaluate d . elabEnv)
  assumeLocal v d_nf $ do
    r <- check r VType
    pure (Sigma v d r, VType)

infer (Rlet v t d b) = do
  t    <- check t VType
  t_nf <- asks (flip evaluate t . elabEnv)

  d    <- check d t_nf
  d_nf <- asks (flip evaluate d . elabEnv)

  defineLocal v t_nf d_nf $ do
    (b, ty) <- infer b
    pure (Let v t d b, ty)

infer Rtype = pure (Type, VType)

infer Rhole = do
  ty <- freshMeta VType
  ty_nf <- asks (flip evaluate ty . elabEnv)
  tm <- freshMeta ty_nf
  pure (tm, ty_nf)

infer (Rlam p v t) = do
  env <- asks elabEnv
  lvl <- asks elabLevel
  dom <- freshMeta VType
  let dom_nf = evaluate env dom
  assumeLocal v dom_nf $ do
    (b, rng) <- infer t
    pure (Lam p v b, VPi p v dom_nf (Closure env (quote (succ lvl) rng)))

infer Rtop  = pure (Top, VType)
infer Runit = pure (Unit, VTop)

infer (Req a b) = do
  t <- freshMeta VType
  t_nf <- asks (flip evaluate t . elabEnv)
  a <- check a t_nf
  b <- check b t_nf
  pure (Id t a b, VType)

infer Rrefl =
  pure (Refl, forAll Im "A" VType $ \a -> forAll Im "x" a $ \x -> VEq a x x)

infer Rcoe =
  pure ( Coe
       , forAll Im "A" VType $ \a ->
         forAll Im "B" VType $ \b ->
         forAll Ex "_" (VEq VType a b) $ \_ ->
         forAll Ex "_" a $ const b
       )

infer Rcong =
  pure ( Cong
       , forAll Im "A" VType $ \a ->
         forAll Im "B" VType $ \b ->
         forAll Im "x" a $ \x ->
         forAll Im "y" a $ \y ->
         forAll Ex "f" (forAll Ex "_" a (const b)) $ \f ->
         forAll Ex "p" (VEq a x y) $ \_ ->
           VEq b (vApp f Ex x) (vApp f Ex y)
       )

infer Rsym =
  pure ( Sym
       , forAll Im "A" VType $ \a -> forAll Im "x" a $ \x -> forAll Im "y" a $ \y -> forAll Ex "p" (VEq a x y) $ \_ -> VEq a y x
       )

infer (Rproj1 e) = do
  (t, ty)   <- infer e
  (_, d, _) <- isSigmaType ty
  pure (Proj1 t, d)

infer (Rproj2 e) = do
  (t, ty)   <- infer e
  t_nf      <- asks (flip evaluate t . elabEnv)
  (_, _, r) <- isSigmaType ty
  pure (Proj2 t, r $$ vProj1 t_nf)

infer c = do
  t <- asks elabSwitches
  when (t >= 128) $
    error $ "Unhandled case in type checker, stack overflew etc: " ++ show c

  t    <- freshMeta VType
  t_nf <- asks (flip evaluate t . elabEnv)

  c <- local (\e -> e { elabSwitches = succ (elabSwitches e)}) $
    check c t_nf

  pure (c, t_nf)

check :: RawExpr -> VTy -> ElabM Term
check (RSrcPos start end expr) ty = local (\st -> st { elabSourcePos = (start, end) }) (check expr ty)

check (Rlam e v t) (VPi e' _ d r) | e == e' = do
  level <- asks (unLvl . elabLevel)
  assumeLocal v d $
    Lam e v <$> check t (r $$ vVar (Bound level))

check t (VPi Im x d r) = do
  level <- asks (unLvl . elabLevel)
  assumeLocal x d $
    Lam Im x <$> check t (r $$ vVar (Bound level))

check (Rlam e v t) ty = do
  (_, d, r) <- isPiType e ty
  level <- asks (unLvl . elabLevel)
  assumeLocal v d $
    Lam e v <$> check t (r $$ vVar (Bound level))

check (Rlet v t d b) ty = do
  t    <- check t VType
  t_nf <- asks (flip evaluate t . elabEnv)

  d    <- check d t_nf
  d_nf <- asks (flip evaluate d . elabEnv)

  defineLocal v t_nf d_nf $ do
    b <- check b ty
    pure (Let v t d b)

check (Rpair a b) ty = do
  (_, d, r) <- isSigmaType ty
  a <- check a d
  a_nf <- asks (flip evaluate a . elabEnv)
  b <- check b (r $$ a_nf)
  pure (Pair a b)

check e ty = do
  (new, e_ty) <- uncurry insert =<< infer e
  unify e_ty ty
    `catchError` \_ -> do
      l <- asks elabLevel
      names <- getNames
      typeError (NotEqual names (quote l (zonk ty)) (quote l (zonk e_ty)))
  pure new

isPiType :: Plicity -> VTy -> ElabM (Text, VTy, Closure)
isPiType i = go . force where
  go (VPi i' a b c)
    | i == i' = pure (a, b, c)
  go ty | not (flexible ty) = do
    l <- asks elabLevel
    names <- getNames
    typeError (NotFunction names (quote l ty))
  go ty = do
    env <- asks elabEnv
    t <- freshMeta VType
    let t_nf = evaluate env t
    assumeLocal "α" t_nf $ do
      r <- freshMeta VType
      unify ty (VPi i "α" t_nf (Closure env r))
      pure ("α", t_nf, Closure env r)

isSigmaType :: VTy -> ElabM (Text, VTy, Closure)
isSigmaType = go . force where
  go (VSigma a b c) = pure (a, b, c)
  go ty = do
    env <- asks elabEnv
    t <- freshMeta VType
    let t_nf = evaluate env t
    assumeLocal "α" t_nf $ do
      r <- freshMeta VType
      unify ty (VSigma "α" t_nf (Closure env r))
      pure ("α", t_nf, Closure env r)