{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedStrings #-}
module Evaluate where

import qualified Control.Exception as Exc
import Control.Monad.Except
import Control.Monad.Reader
import Control.Concurrent

import qualified Data.IntMap.Strict as IntMap
import qualified Data.Sequence as Seq
import qualified Data.Text as T

import Elaboration.Monad

import GHC.Stack (HasCallStack)

import Generics.SYB (mkT, everywhere)

import Syntax

import System.IO.Unsafe

import Value
import Data.Foldable
import Syntax.Pretty (showWithPrec)

evaluate :: HasCallStack => Env -> Term -> Value
evaluate env (Var (Bound i)) =
  case Seq.lookup i (locals env) of
    Just x -> x
    Nothing -> error $ "Variable of index " ++ show i ++ " not in scope"

evaluate _ (Con t) = VGlued (HCon t) mempty Nothing

evaluate _   Type = VType

evaluate env (Pi p t d r)  = VPi p t (evaluate env d) (Closure env r)
evaluate env (Lam p t b)   = VLam p t (Closure env b)
evaluate env (App p f x)   = vApp (evaluate env f) p (evaluate env x)

evaluate env (Sigma t d r) = VSigma t (evaluate env d) (Closure env r)
evaluate env (Pair a b)    = VPair (evaluate env a) (evaluate env b)
evaluate env (Proj1 a)     = vProj1 (evaluate env a)
evaluate env (Proj2 a)     = vProj2 (evaluate env a)

evaluate _   (Meta m)         = VGlued (HMeta m) mempty Nothing
evaluate env (NewMeta m mask) = VGlued (HMeta m) (getVals (locals env) mask) Nothing where
  getVals Seq.Empty []                    = Seq.Empty
  getVals (v Seq.:<| seq) (BDBound _:bds)   = AppEx v Seq.:<| getVals seq bds
  getVals (_ Seq.:<| seq) (BDDefined _:bds) = getVals seq bds

evaluate _ Top = VTop
evaluate _ Unit = VUnit

evaluate _ Refl = VGlued HRefl mempty Nothing
evaluate _ Coe  =
  function Im (T.pack "A") $ \a ->
  function Im (T.pack "B") $ \b ->
  function Ex (T.pack "p") $ \p ->
  function Ex (T.pack "x") $ \x ->
    vCoe a b p x

evaluate _ Cong =
  function Im (T.pack "A") $ \a ->
  function Im (T.pack "B") $ \b ->
  function Im (T.pack "x") $ \x ->
  function Im (T.pack "y") $ \y ->
  function Ex (T.pack "f") $ \f ->
  function Ex (T.pack "p") $ \p ->
    vCong a b x y f p

evaluate _ Sym =
  function Im (T.pack "A") $ \a ->
  function Im (T.pack "x") $ \x ->
  function Im (T.pack "y") $ \y ->
  function Ex (T.pack "p") $ \p ->
    vSym a x y p

evaluate e (Let _ _ c d) = evaluate e' d where
  e' = e { locals = evaluate e c Seq.:<| locals e }

evaluate env (Id a b c) = vId (evaluate env a) (evaluate env b) (evaluate env c)

vId :: Value -> Value -> Value -> Value
vId kind a b =
  let stuck  = VEq kind a b
      solve  = VEqG kind a b
      never  = solve vBottom
      always = solve VTop
   in
  case force kind of
    VType ->
      case (a, b) of
        (VTop, VTop) -> always

        (VType, VType) -> always

        (VEqG _ _ _ a, b) -> vId VType a b
        (a, VEqG _ _ _ b) -> vId VType a b

        (VEq a _ _, VEq b _ _) -> vId VType a b

        (VPi i _ d r, VPi i' _ d' r')
          | i == i' ->
            solve $
              exists "p" (vId VType d d') $ \p ->
              forAll Ex "x" d $ \x ->
                vId VType (r $$ x) (r' $$ vCoe d d' p x)
          | otherwise -> never

        (VSigma _ d r, VSigma _ d' r') ->
          solve $
            exists "p" (vId VType d d') $ \p ->
            forAll Ex "x" d $ \x ->
              vId VType (r $$ x) (r' $$ vCoe d d' p x)

        (VNe _ _, _) -> stuck
        (_, VNe _ _) -> stuck

        _ -> never

    VTop -> always

    VPi i t dom cod ->
      solve $ forAll i t dom \vl -> vId (cod $$ vl) (vApp a i vl) (vApp b i vl)

    VSigma t dom cod ->
      -- a = (x, p)
      -- b = (y, q)
      -- (a, b) ≡ (c, d) : (x : A) * P x
      --   ~> (path : a == c) * coe (cong A Type P path) b == d
      let x = vProj1 a
          y = vProj1 b
          p = vProj2 a
          q = vProj2 b
       in solve $
         exists t (vId dom x y) $ \pr ->
           vId (cod $$ y) (vCoe (cod $$ x) (cod $$ y) (vCong dom VType x y (function Ex (T.pack "x") (cod $$)) pr) p) q

    VEq{} -> solve VTop

    _ -> stuck

vBottom :: Value
vBottom = forAll Im "A" VType id

vCoe :: VTy -> VTy -> Value -> Value -> Value
-- vCoe _ _ (VGlued HRefl _ _) element = element
vCoe (VPi i _ d r) ty2@(VPi i' _ d' r') p f
  | i /= i' = vApp p Ex ty2
  | otherwise =
    function i "x" $ \x ->
      let p1 = vProj1 p -- d == d'
          p2 = vProj2 p -- (x : A) -> r x == r' (coe p1 x)
          x0 = vCoe d' d (vSym VType d d' p1) x
       in vCoe (r $$ x0) (r' $$ x) (vApp p2 Ex x0) (vApp f i x0)
 
vCoe tyA tyB proof element = splitFibration tyA tyB $ VGlued HCoe (Seq.fromList spine) Nothing where
  spine = [AppIm tyA, AppIm tyB, AppEx proof, AppEx element]

  -- Types are split fibrations
  -- coe {A} {A} p x = x even when p /= refl
  splitFibration tA tB vstuck = unsafeDupablePerformIO $ do
    old <- readMVar elabMetas
    t <- runElab (unify tA tB) emptyElabState
    case t of
      Left _ -> do
        swapMVar elabMetas old
        pure vstuck
      Right _ -> pure element

vCong :: Value -> Value -> Value -> Value -> Value -> Value -> Value
vCong _a c _x _y g (force -> VGlued HCong (toList -> [AppIm a, AppIm _b, AppIm x, AppIm y, AppEx f, AppEx p]) _) =
  VGlued HCong (Seq.fromList [AppIm a, AppIm c, AppIm x, AppIm y, AppEx (function Ex "x" (vApp g Ex . vApp f Ex)), AppEx p]) Nothing
vCong _a b _x _ f (force -> VGlued HRefl (toList -> [AppIm _, AppIm x]) _) =
  VGlued HRefl (Seq.fromList [AppIm b, AppIm (vApp f Ex x)]) Nothing
vCong a b x y f p =
  VGlued HCong (Seq.fromList [AppIm a, AppIm b, AppIm x, AppIm y, AppEx f, AppEx p]) Nothing

vSym :: Value -> Value -> Value -> Value -> Value
vSym a _ y (VGlued HRefl _ Nothing) =
  VGlued HRefl (Seq.fromList [AppIm a, AppIm y]) Nothing
vSym _ _ _ (VGlued HSym (toList -> [_a, _y, _x, AppEx p]) Nothing) = p
vSym a x y p = VGlued HSym (Seq.fromList [AppIm a, AppIm x, AppIm y, AppEx p]) Nothing

vApp :: HasCallStack => Value -> Plicity -> Value -> Value
vApp (VGlued x s v) p r = VGlued x (s Seq.|> thing) (fmap (\v -> vApp v p r) v) where
  thing =
    case p of
      Ex -> AppEx r
      Im -> AppIm r
vApp (VLam _ _ c) _ a = c $$ a
vApp _fun _plic _arg = error "invalid application"

vProj1 :: Value -> Value
vProj1 (VPair a _) = a
vProj1 x           = VProj1 x

vProj2 :: Value -> Value
vProj2 (VPair _ a) = a
vProj2 x           = VProj2 x

($$) :: HasCallStack => Closure -> Value -> Value
Closure e t $$ v = evaluate e' t where
  e' = e { locals = v Seq.:<| locals e }
ClMeta (MetaFun f) $$ v = f v

forAll :: Plicity -> T.Text -> Value -> (Value -> Value) -> Value
forAll i t d = VPi i t d . ClMeta . MetaFun

exists :: T.Text -> Value -> (Value -> Value) -> Value
exists t d = VSigma t d . ClMeta . MetaFun

function :: Plicity -> T.Text -> (Value -> Value) -> Value
function i t = VLam i t . ClMeta . MetaFun

quote :: HasCallStack => Level -> Value -> Term
quote _ VType = Type

quote l (VPi p t d r) = Pi p t (quote l d) (quote (succ l) (r $$ vVar (Bound (unLvl l))))
quote l (VLam p t  b) = Lam p t (quote (succ l) (b $$ vVar (Bound (unLvl l))))

quote l (VSigma t d r) = Sigma t (quote l d) (quote (succ l) (r $$ vVar (Bound (unLvl l))))
quote l (VPair a b)    = Pair (quote l a) (quote l b)
quote l (VProj1 a)     = Proj1 (quote l a)
quote l (VProj2 a)     = Proj2 (quote l a)

quote _ VTop  = Top
quote _ VUnit = Unit

quote l (VEq a b c) = Id (quote l a) (quote l b) (quote l c)
-- quote l (VEqG _ _ _ d) = quote l d
quote l (VEqG a b c _) = Id (quote l a) (quote l b) (quote l c)

quote l (VGlued v s _) = foldl app v' s where
  v' = case v of
    HVar (Bound i) -> Bv (lvl2Ix l (Lvl i))
    HCon t         -> Con t
    HMeta m        -> Meta m
    HRefl          -> Refl
    HCoe           -> Coe
    HCong          -> Cong
    HSym           -> Sym

  app f (AppEx t) = App Ex f (quote l t)
  app f (AppIm t) = App Im f (quote l t)
  app f SProj1 = Proj1 f
  app f SProj2 = Proj2 f

force :: Value -> Value
force = unsafeDupablePerformIO . go where
  go stuck@(VGlued (HMeta (MV m)) sp Nothing) = do
    t <- readMVar elabMetas
    case IntMap.lookup m t of
      Just (Solved vl) -> go $ foldl vAppSp vl sp
      _ -> pure stuck
  go (VGlued _ _ (Just vl)) = go vl
  go x = pure x

vAppSp :: Value -> SpineThing -> Value
vAppSp vl (AppEx f) = vApp vl Ex f
vAppSp vl (AppIm f) = vApp vl Im f
vAppSp vl SProj1 = vProj1 vl
vAppSp vl SProj2 = vProj2 vl

zonk :: Value -> Value
zonk (VLam vis var body) = VLam vis var (ClMeta (MetaFun (\v -> zonk (body $$ v))))
zonk (VPi vis var dom body) = VPi vis var (zonk dom) (ClMeta (MetaFun (\v -> zonk (body $$ v))))
zonk (VSigma var dom body) = VSigma var (zonk dom) (ClMeta (MetaFun (\v -> zonk (body $$ v))))
zonk t = everywhere (mkT force) t

unify :: VTy -> VTy -> ElabM ()
unify a b = asks elabLevel >>= flip go (a, b) where
  go, go' :: Level -> (VTy, VTy) -> ElabM ()
  go' l (VGlued h sp x, VGlued h' sp' y)
    | h == h'                  = goSpine l sp sp'
    | Just x <- x, Just y <- y = go l (x, y)

  -- flexible head (solve meta)
  go' l (VGlued (HMeta m) sp _, y) = solveMeta m sp y
  go' _ (x, VGlued (HMeta m) sp _) = solveMeta m sp x

  -- rigid heads (compare unfolding)
  go' l (VGlued _ _ (Just x), y) = go l (x, y)
  go' l (x, VGlued _ _ (Just y)) = go l (x, y)

  go' _ (VType, VType) = pure ()
  go' _ (VTop, VTop) = pure ()
  go' _ (VUnit, VUnit) = pure ()

  go' l (VPi i _ d r, VPi i' _ d' r') | i == i' = do
    go l (d, d')
    let i = unLvl l
    go (succ l) (r $$ vVar (Bound i), r' $$ vVar (Bound i))

  go' l (VSigma _ d r, VSigma _ d' r') = do
    go l (d, d')
    let i = unLvl l
    go (succ l) (r $$ vVar (Bound i), r' $$ vVar (Bound i))

  go' l (VLam i _ r, VLam i' _ r') | i == i' = do
    let i = unLvl l
    go (succ l) (r $$ vVar (Bound i), r' $$ vVar (Bound i))

  go' l (VLam p _ r, t) = do
    let i = unLvl l
    go (succ l) (r $$ vVar (Bound i), vApp t p (vVar (Bound i)))

  go' l (r, VLam p _ t) = do
    let i = unLvl l
    go (succ l) (vApp r p (vVar (Bound i)), t $$ vVar (Bound i))

  go' l (VEqG a b c _, VEqG a' b' c' _) = go l (a, a') *> go l (b, b') *> go l (c, c')
  go' l (VEq a b c, VEqG a' b' c' _)    = go l (a, a') *> go l (b, b') *> go l (c, c')
  go' l (VEqG a b c _, VEq a' b' c')    = go l (a, a') *> go l (b, b') *> go l (c, c')
  go' l (VEq a b c, VEq a' b' c')       = go l (a, a') *> go l (b, b') *> go l (c, c')

  go' l (VEqG _ _ _ a, b) = go l (a, b)
  go' l (a, VEqG _ _ _ b) = go l (a, b)

  go' l (VProj1 a, VProj1 b) = go l (a, b)
  go' l (VProj2 a, VProj2 b) = go l (a, b)
  go' l (VPair a b, VPair a' b') = go l (a, a') >> go l (b, b')

  go' l (a, b) = do
    ns <- getNames
    typeError (NotEqual ns (quote l a) (quote l b))
  
  go l (a, b) = go' l (force a, force b)

  goSpine _ Seq.Empty Seq.Empty           = pure ()
  goSpine l (AppEx x Seq.:<| xs) (AppEx y Seq.:<| ys) = do
    go l (x, y)
    goSpine l xs ys
  goSpine l (AppIm x Seq.:<| xs) (AppIm y Seq.:<| ys) = do
    go l (x, y)
    goSpine l xs ys
  goSpine l (x Seq.:<| xs) (y Seq.:<| ys) | x == y = goSpine l xs ys
  goSpine l _ _ = do
    ns <- getNames
    typeError (NotEqual ns (quote l a) (quote l b))

solveMeta :: HasCallStack => MetaVar -> Seq.Seq SpineThing -> VTy -> ElabM ()
solveMeta (MV meta) spine rhs =
  do
    level <- asks elabLevel

    pren <- invert level spine
    rhs <- rename (MV meta) pren rhs
    let solution = evaluate emptyEnv (lams (dom pren) rhs) 

    -- need to deepSeq solutions here
    -- no deepSeq? no problem

    liftIO $ Exc.evaluate (length (show solution))

    liftIO . modifyMVar_ elabMetas $ pure . IntMap.insert meta (Solved solution)
  `catchError` \case
    [] -> do
      level <- asks elabLevel
      names <- getNames
      typeError (CantSolveMeta names (quote level (VGlued (HMeta (MV meta)) spine Nothing)) (quote level rhs))
    cs -> throwError cs

elabMetas :: MVar (IntMap.IntMap Meta)
elabMetas = unsafeDupablePerformIO (newMVar mempty)
{-# NOINLINE elabMetas #-}

lams :: Level -> Term -> Term
lams l = go (Lvl 0) where
  go x t | x == l = t
  go x t = Lam Ex (T.pack ("x" ++ show (unLvl x))) $ go (succ x) t

data PartialRen
  = PRen { dom :: {-# UNPACK #-} !Level
         , rng :: {-# UNPACK #-} !Level
         , sub :: IntMap.IntMap Level
         }
  deriving (Eq, Show, Ord)

liftRen :: PartialRen -> PartialRen
liftRen (PRen d r s) = PRen (succ d) (succ r) (IntMap.insert (unLvl r) d s)

invert :: Level -> Seq.Seq SpineThing -> ElabM PartialRen
invert gamma spine =
  do
    (dom, ren) <- go spine
    pure (PRen dom gamma ren)
  where
    go Seq.Empty = pure (Lvl 0, mempty)
    go (sp Seq.:|> AppEx t) = do
      (dom, ren) <- go sp
      case force t of
        VGlued (HVar (Bound l)) Seq.Empty _
          | IntMap.notMember l ren -> pure (succ dom, IntMap.insert l dom ren)
        _ -> throwError []
    go (_ Seq.:|> _) = throwError []

rename :: HasCallStack => MetaVar -> PartialRen -> Value -> ElabM Term
rename meta pren = go pren where
  go :: HasCallStack => PartialRen -> Value -> ElabM Term
  go pren (VGlued (HMeta m) sp _)
    | m == meta = throwError []
    | otherwise = goSp pren (Meta m) sp

  go pren (VGlued (HVar (Bound m)) sp _) =
    case IntMap.lookup m (sub pren) of
      Just v -> goSp pren (Bv (lvl2Ix (dom pren) v)) sp
      Nothing -> throwError []

  go pren (VGlued h sp _) = goHead h >>= \h -> goSp pren h sp where
    goHead HRefl      = pure Refl
    goHead HCong      = pure Cong
    goHead HCoe       = pure Coe
    goHead HSym       = pure Sym

  go pren (VPi p t d r) = Pi p t <$> go pren d <*> go (liftRen pren) (r $$ vVar (Bound (unLvl (rng pren))))
  go pren (VLam p t x)  = Lam p t <$> go (liftRen pren) (x $$ vVar (Bound (unLvl (rng pren))))
  go _ VType            = pure Type
  go _ VTop             = pure Top
  go _ VUnit            = pure Unit
  
  go pren (VSigma t d r) = Sigma t <$> go pren d <*> go (liftRen pren) (r $$ vVar (Bound (unLvl (rng pren))))
  go pren (VPair a b) = Pair <$> go pren a <*> go pren b
  go pren (VProj1 a) = Proj1 <$> go pren a
  go pren (VProj2 a) = Proj2 <$> go pren a

  go pren (VEq a b c)    = Id <$> go pren a <*> go pren b <*> go pren c
  go pren (VEqG _ _ _ d) = go pren d

  -- go pren x = error (show x)

  goSp _ t Seq.Empty               = pure t
  goSp pren t (sp Seq.:|> AppEx tm) = App Ex <$> goSp pren t sp <*> go pren tm
  goSp pren t (sp Seq.:|> AppIm tm) = App Im <$> goSp pren t sp <*> go pren tm
  goSp pren t (sp Seq.:|> SProj1) = Proj1 <$> goSp pren t sp
  goSp pren t (sp Seq.:|> SProj2) = Proj2 <$> goSp pren t sp