{-# 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.Sequence as Seq
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Sequence (Seq)
import Data.Traversable
import Data.Set (Set)
import Data.Typeable
import Data.Foldable
import Data.IORef
import Data.Maybe

import Elab.Monad

import Presyntax.Presyntax (Plicity(..))

import Syntax.Pretty
import Syntax

import System.IO.Unsafe

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

eval :: Term -> ElabM Value
eval t = asks (flip eval' 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 args
    Nothing -> pure vl
forceIO x = pure x

applProj :: Value -> Projection -> Value
applProj fun (PApp p arg)     = vApp p fun arg
applProj fun (PIElim l x y i) = ielim l x y fun i
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 sp'
        Nothing -> pure $ VNe hd sp'
    hd -> pure $ VNe hd sp'
  where
    zonkSp (PApp p x) = PApp p <$> zonkIO x
    zonkSp (PIElim l x y i) = PIElim <$> zonkIO l <*> zonkIO x <*> zonkIO y <*> zonkIO i
    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

zonkIO (VPath line x y) = VPath <$> zonkIO line <*> zonkIO x <*> zonkIO y
zonkIO (VLine line f) = VLine <$> zonkIO line <*> zonkIO f

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

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

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

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

eval' _ (Meta m) = VNe (HMeta m) mempty

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

eval' e (Pair a b) = VPair (eval' e a) (eval' e b)

eval' e (Proj1 a) = vProj1 (eval' e a)
eval' e (Proj2 a) = vProj2 (eval' e a)

eval' _ Type = VType
eval' _ Typeω = VTypeω
eval' _ I = VI
eval' _ I0 = VI0
eval' _ I1 = VI1

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

eval' e (PathP l a b) = VPath (eval' e l) (eval' e a) (eval' e b)
eval' e (IElim l x y f i) = ielim (eval' e l) (eval' e x) (eval' e y) (eval' e f) (eval' e i)
eval' e (PathIntro p f) = VLine (eval' e p) (eval' e f)

vApp :: Plicity -> Value -> Value -> Value
vApp p (VLam p' k) arg
  | p == p'   = clCont k arg
  | otherwise = error $ "wrong plicity " ++ show p ++ " vs " ++ show p' ++ " in app " ++ show (App p (quote (VLam p' k)) (quote arg))
vApp p (VNe h sp) arg  = VNe h (sp Seq.:|> PApp p arg)
vApp _ x _             = error $ "can't apply " ++ show x

(@@) :: Value -> Value -> Value
(@@) = vApp Ex
infixl 9 @@

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

vProj2 :: Value -> Value
vProj2 (VPair _ b) = b
vProj2 (VNe h sp)  = VNe h (sp Seq.:|> PProj2)
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) (Seq.zip a a')
  
  go (VNe _hd (_ Seq.:|> PIElim _l x y i)) rhs =
    case force i of
      VI0 -> unify' x rhs
      VI1 -> unify' y rhs
      _   -> 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 (VINot x)   (VINot y)     = unify' x y
  go (VIAnd x y) (VIAnd x' y') = unify' x x' *> unify' y y'
  go (VIOr x y)  (VIOr x' y')  = unify' x x' *> unify' y y'

  go (VPath l x y) (VPath l' x' y') = do
    unify' l l'
    unify' x x'
    unify' y y'
  
  go (VLine l p) p' = do
    n <- VVar . Bound <$> newName
    unify (p @@ n) (ielim l (l @@ VI0) (l @@ VI1) p' n)

  go p' (VLine l p) = do
    n <- VVar . Bound <$> newName
    unify (ielim l (l @@ VI0) (l @@ VI1) p' n) (p @@ n)

  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 (PIElim _ _ _ i) (PIElim _ _ _ j) = unify' i j

  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 Im f (quote meta)))
VType `isConvertibleTo` VTypeω = pure id

VPi p d (Closure _ k) `isConvertibleTo` VPi p' d' (Closure _ k') | p == p' = do
  wp <- d' `isConvertibleTo` d
  n <- newName
  wp_n <- eval (Lam Ex n (wp (Ref (Bound n))))

  wp' <- k (VVar (Bound n)) `isConvertibleTo` k' (wp_n @@ VVar (Bound n))
  pure (\f -> Lam p n (wp' (App p f (wp (Ref (Bound n))))))

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) (Seq.fromList (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 -> Seq Projection -> Value -> ElabM ()
solveMeta m@(MV _ cell) sp rhs = do
  env <- ask
  names <- checkSpine Set.empty sp
  checkScope (Set.fromList (Bound <$> names)) rhs
    `withNote` hsep [prettyTm (quote (VNe (HMeta m) sp)), pretty '≡', prettyTm (quote rhs)]
  let tm = quote rhs
      lam = eval' 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 (PIElim l x y i) = traverse_ (checkScope scope) [l, x, y, i]
    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

checkScope s (VPath line a b) = traverse_ (checkScope s) [line, a, b]
checkScope s (VLine _ line) = checkScope s line

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

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

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