{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Elab.Eval where

import Control.Monad.Reader
import Control.Exception

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Traversable
import Data.Set (Set)
import Data.Typeable
import Data.Foldable
import Data.IORef
import Data.Maybe

import {-# SOURCE #-} Elab.WiredIn
import Elab.Monad

import Presyntax.Presyntax (Plicity(..))

import Syntax

import System.IO.Unsafe
import Syntax.Pretty

eval :: Term -> ElabM Value
eval t = asks (flip evalWithEnv t)

forceIO :: MonadIO m => Value -> m Value
forceIO vl@(VNe (HMeta (MV _ cell)) args) = do
  solved <- liftIO $ readIORef cell
  case solved of
    Just vl -> forceIO $ foldl applProj vl (reverse args)
    Nothing -> pure vl
forceIO x = pure x

applProj :: Value -> Projection -> Value
applProj fun (PApp p arg) = vApp p fun arg
applProj fun PProj1 = vProj1 fun
applProj fun PProj2 = vProj2 fun

force :: Value -> Value
force = unsafePerformIO . forceIO

-- everywhere force
zonkIO :: Value -> IO Value
zonkIO (VNe hd sp) = do
  sp' <- traverse zonkSp sp
  case hd of
    HMeta (MV _ cell) -> do
      solved <- liftIO $ readIORef cell
      case solved of
        Just vl -> zonkIO $ foldl applProj vl (reverse sp')
        Nothing -> pure $ VNe hd sp'
    hd -> pure $ VNe hd sp'
  where
    zonkSp (PApp p x) = PApp p <$> zonkIO x
    zonkSp PProj1 = pure PProj1
    zonkSp PProj2 = pure PProj2
zonkIO (VLam p (Closure s k)) = pure $ VLam p (Closure s (zonk . k))
zonkIO (VPi p d (Closure s k)) = VPi p <$> zonkIO d <*> pure (Closure s (zonk . k))
zonkIO (VSigma d (Closure s k)) = VSigma <$> zonkIO d <*> pure (Closure s (zonk . k))
zonkIO (VPair a b) = VPair <$> zonkIO a <*> zonkIO b

-- Sorts
zonkIO VType = pure VType
zonkIO VTypeω = pure VTypeω

zonkIO VI = pure VI
zonkIO VI0 = pure VI0
zonkIO VI1 = pure VI1

zonkIO (VIAnd x y) = iand <$> zonkIO x <*> zonkIO y
zonkIO (VIOr x y) = ior <$> zonkIO x <*> zonkIO y
zonkIO (VINot x) = inot <$> zonkIO x

zonk :: Value -> Value
zonk = unsafePerformIO . zonkIO

evalWithEnv :: ElabEnv -> Term -> Value
evalWithEnv env (Ref x) =
  case Map.lookup x (getEnv env) of
    Just (_, vl) -> vl
    _ -> error "variable not in scope when evaluating"
evalWithEnv env (App p f x) = vApp p (evalWithEnv env f) (evalWithEnv env x)

evalWithEnv env (Lam p s t) =
  VLam p $ Closure s $ \a ->
    evalWithEnv env { getEnv = Map.insert (Bound s) (error "type of abs", a) (getEnv env) } t

evalWithEnv env (Pi p s d t) =
  VPi p (evalWithEnv env d) $ Closure s $ \a ->
    evalWithEnv env { getEnv = (Map.insert (Bound s) (error "type of abs", a) (getEnv env))} t

evalWithEnv _ (Meta m) = VNe (HMeta m) []

evalWithEnv env (Sigma s d t) =
  VSigma (evalWithEnv env d) $ Closure s $ \a ->
    evalWithEnv env { getEnv = Map.insert (Bound s) (error "type of abs", a) (getEnv env) } t

evalWithEnv e (Pair a b) = VPair (evalWithEnv e a) (evalWithEnv e b) 

evalWithEnv e (Proj1 a) = vProj1 (evalWithEnv e a)
evalWithEnv e (Proj2 a) = vProj2 (evalWithEnv e a)

evalWithEnv _ Type = VType
evalWithEnv _ Typeω = VTypeω
evalWithEnv _ I = VI
evalWithEnv _ I0 = VI0
evalWithEnv _ I1 = VI1

evalWithEnv e (IAnd x y) = iand (evalWithEnv e x) (evalWithEnv e y)
evalWithEnv e (IOr x y) = ior (evalWithEnv e x) (evalWithEnv e y)
evalWithEnv e (INot x) = inot (evalWithEnv e x)

vApp :: Plicity -> Value -> Value -> Value
vApp p (VLam p' k) arg = assert (p == p') $ clCont k arg
vApp p (VNe h sp) arg  = VNe h (PApp p arg:sp)
vApp _ x _             = error $ "can't apply " ++ show x

vProj1 :: Value -> Value
vProj1 (VPair a _) = a
vProj1 (VNe h sp) = VNe h (PProj1:sp)
vProj1 x = error $ "can't proj1 " ++ show x

vProj2 :: Value -> Value
vProj2 (VPair _ b) = b
vProj2 (VNe h sp)  = VNe h (PProj2:sp)
vProj2 x = error $ "can't proj2 " ++ show x

data NotEqual = NotEqual Value Value
  deriving (Show, Typeable, Exception)

unify' :: Value -> Value -> ElabM ()
unify' topa topb = join $ go <$> forceIO topa <*> forceIO topb where
  go (VNe (HMeta mv) sp) rhs = solveMeta mv sp rhs
  go rhs (VNe (HMeta mv) sp) = solveMeta mv sp rhs

  go (VNe x a) (VNe x' a')
    | x == x', length a == length a' =
      traverse_ (uncurry unify'Spine) (zip a a')
    | otherwise = fail

  go (VLam p (Closure _ k)) vl = do
    t <- VVar . Bound <$> newName
    unify' (k t) (vApp p vl t)

  go vl (VLam p (Closure _ k)) = do
    t <- VVar . Bound <$> newName
    unify' (vApp p vl t) (k t)

  go (VPair a b) vl = unify' a (vProj1 vl) *> unify' b (vProj2 vl)
  go vl (VPair a b) = unify' (vProj1 vl) a *> unify' (vProj2 vl) b

  go (VPi p d (Closure _ k)) (VPi p' d' (Closure _ k')) | p == p' = do
    t <- VVar . Bound <$> newName
    unify' d d'
    unify' (k t) (k' t)

  go (VSigma d (Closure _ k)) (VSigma d' (Closure _ k')) = do
    t <- VVar . Bound <$> newName
    unify' d d'
    unify' (k t) (k' t)

  go VType VType   = pure ()
  go VTypeω VTypeω = pure ()

  go VI VI   = pure ()
  go VI0 VI0 = pure ()
  go VI1 VI1 = pure ()

  go _ _ = fail

  fail = liftIO . throwIO $ NotEqual topa topb

  unify'Spine (PApp a v) (PApp a' v')
    | a == a' = unify' v v'

  unify'Spine PProj1 PProj1 = pure ()
  unify'Spine PProj2 PProj2 = pure ()

  unify'Spine _ _ = fail

unify :: Value -> Value -> ElabM ()
unify a b = unify' a b `catchElab` \(_ :: SomeException) -> liftIO $ throwIO (NotEqual a b)

isConvertibleTo :: Value -> Value -> ElabM (Term -> Term)
VPi Im d (Closure _v k) `isConvertibleTo` ty = do
  meta <- newMeta d
  cont <- k meta `isConvertibleTo` ty
  pure (\f -> cont (App Ex f (quote meta)))
VType `isConvertibleTo` VTypeω = pure id
isConvertibleTo a b = do
  unify' a b
  pure id

newMeta :: Value -> ElabM Value
newMeta _dom = do
  n <- newName
  c <- liftIO $ newIORef Nothing
  let m = MV n c

  env <- asks getEnv

  t <- for (Map.toList env) $ \(n, _) -> pure $
    case n of
      Bound{} -> Just (PApp Ex (VVar n))
      _ -> Nothing

  pure (VNe (HMeta m) (catMaybes t))

newName :: MonadIO m => m T.Text
newName = liftIO $ do
  x <- atomicModifyIORef _nameCounter $ \x -> (x + 1, x + 1)
  pure (T.pack (show x))

_nameCounter :: IORef Int
_nameCounter = unsafePerformIO $ newIORef 0
{-# NOINLINE _nameCounter #-}

solveMeta :: MV -> [Projection] -> Value -> ElabM ()
solveMeta m@(MV _ cell) sp rhs = do
  env <- ask
  liftIO $ putStrLn (showValue (VNe (HMeta m) sp) ++ " ≡? " ++ showValue rhs)
  names <- checkSpine Set.empty sp
  checkScope (Set.fromList (Bound <$> names)) rhs
  let tm = quote rhs
      lam = evalWithEnv env $ foldr (Lam Ex) tm names
  liftIO . atomicModifyIORef' cell $ \case
    Just _ -> error "filled cell in solvedMeta"
    Nothing -> (Just lam, ())

checkScope :: Set Name -> Value -> ElabM ()
checkScope scope (VNe h sp) =
  do
    case h of
      HVar v@Bound{} ->
        unless (v `Set.member` scope) . liftIO . throwIO $
          NotInScope v
      HVar{} -> pure ()
      HMeta{} -> pure ()
    traverse_ checkProj sp
  where
    checkProj (PApp _ t) = checkScope scope t
    checkProj PProj1 = pure ()
    checkProj PProj2 = pure ()

checkScope scope (VLam _ (Closure n k)) =
  checkScope (Set.insert (Bound n) scope) (k (VVar (Bound n)))

checkScope scope (VPi _ d (Closure n k)) = do
  checkScope scope d
  checkScope (Set.insert (Bound n) scope) (k (VVar (Bound n)))

checkScope scope (VSigma d (Closure n k)) = do
  checkScope scope d
  checkScope (Set.insert (Bound n) scope) (k (VVar (Bound n)))

checkScope s (VPair a b) = traverse_ (checkScope s) [a, b]

checkScope _ VType = pure ()
checkScope _ VTypeω = pure ()

checkScope _ VI = pure ()
checkScope _ VI0 = pure ()
checkScope _ VI1 = pure ()

checkScope s (VIAnd x y) = traverse_ (checkScope s) [x, y]
checkScope s (VIOr x y)  = traverse_ (checkScope s) [x, y]
checkScope s (VINot x)   = checkScope s x

checkSpine :: Set Name -> [Projection] -> ElabM [T.Text]
checkSpine scope (PApp Ex (VVar n@(Bound t)):xs)
  | n `Set.member` scope = liftIO . throwIO $ NonLinearSpine n
  | otherwise = (t:) <$> checkSpine scope xs
checkSpine _     (p:_) = liftIO . throwIO $ SpineProj p
checkSpine _     [] = pure []

newtype NonLinearSpine = NonLinearSpine { getDupeName :: Name }
  deriving (Show, Typeable, Exception)

newtype SpineProjection = SpineProj { getSpineProjection :: Projection }
  deriving (Show, Typeable, Exception)