module Quox.Equal

import public Quox.Syntax
import public Quox.Definition
import public Quox.Typing
import public Control.Monad.Either
import public Control.Monad.Reader
import Data.Maybe


private %inline
ClashE : EqMode -> Elim q d n -> Elim q d n -> Error q
ClashE mode = ClashT mode `on` E


public export
record Env' q (isGlobal : q -> Type) where
  constructor MakeEnv
  defs : Definitions' q isGlobal
  mode : EqMode

public export
0 Env : (q : Type) -> IsQty q => Type
Env q = Env' q IsGlobal

public export
0 HasEnv' : (q : Type) -> (q -> Type) -> (Type -> Type) -> Type
HasEnv' q isGlobal = MonadReader (Env' q isGlobal)

public export
0 HasEnv : (q : Type) -> IsQty q => (Type -> Type) -> Type
HasEnv q = HasEnv' q IsGlobal


public export
0 CanEqual' : (q : Type) -> (q -> Type) -> (Type -> Type) -> Type
CanEqual' q isGlobal m = (HasErr q m, HasEnv' q isGlobal m)

public export
0 CanEqual : (q : Type) -> IsQty q => (Type -> Type) -> Type
CanEqual q = CanEqual' q IsGlobal



private %inline
mode : HasEnv' _ _ m => m EqMode
mode = asks mode

private %inline
clashT : CanEqual' q _ m => Term q d n -> Term q d n -> m a
clashT s t = throwError $ ClashT !mode s t

private %inline
clashE : CanEqual' q _ m => Elim q d n -> Elim q d n -> m a
clashE e f = throwError $ ClashE !mode e f

private %inline
defE : HasEnv' q _ m => Name -> m (Maybe (Elim q d n))
defE x = asks $ \env => do
  g <- lookup x env.defs
  pure $ (!g.term).get :# g.type.get

private %inline
defT : HasEnv' q _ m => Name -> m (Maybe (Term q d n))
defT x = map E <$> defE x

export %inline
compareU' : HasEnv' q _ m => Universe -> Universe -> m Bool
compareU' i j = pure $
  case !mode of Equal => i == j; Sub => i <= j

export %inline
compareU : CanEqual' q _ m => Universe -> Universe -> m ()
compareU k l = unless !(compareU' k l) $
  throwError $ ClashU !mode k l

mutual
  private covering
  compareTN' : CanEqual' q _ m => Eq q =>
               (s, t : Term q 0 n) ->
               (0 _ : NotRedexT s) -> (0 _ : NotRedexT t) -> m ()

  compareTN' (E e) (E f) ps pt = compareE0 e f
  -- if either term is a def, try to unfold it
  compareTN' s@(E (F x)) t ps pt = do
    Just s' <- defT x  | Nothing => clashT s t
    compareT0 s' t
  compareTN' s t@(E (F y)) ps pt = do
    Just t' <- defT y  | Nothing => clashT s t
    compareT0 s t'
  compareTN' s@(E _) t _ _ = clashT s t

  compareTN' (TYPE k) (TYPE l) _ _ = compareU k l
  compareTN' s@(TYPE _) t _ _ = clashT s t

  compareTN' (Pi qty1 _ arg1 res1) (Pi qty2 _ arg2 res2) _ _ = do
    unless (qty1 == qty2) $ throwError $ ClashQ qty1 qty2
    compareT0  arg2 arg1 -- reversed for contravariant domain
    compareTS0 res1 res2
  compareTN' s@(Pi {}) t _ _ = clashT s t

  -- [todo] eta
  compareTN' (Lam _ body1) (Lam _ body2) _ _ =
    local {mode := Equal} $ compareTS0 body1 body2
  compareTN' s@(Lam {}) t _ _ = clashT s t

  compareTN' (CloT  {}) _ ps _ = void $ ps IsCloT
  compareTN' (DCloT {}) _ ps _ = void $ ps IsDCloT

  private covering
  compareEN' : CanEqual' q _ m => Eq q =>
               (e, f : Elim q 0 n) ->
               (0 _ : NotRedexE e) -> (0 _ : NotRedexE f) -> m ()

  compareEN' e@(F x) f@(F y) _ _ =
    if x == y then pure () else
    case (!(defE x), !(defE y)) of
         (Nothing, Nothing) => clashE e f
         (s', t') => compareE0 (fromMaybe e s') (fromMaybe f t')
  compareEN' e@(F x) f _ _ = do
    Just e' <- defE x  | Nothing => clashE e f
    compareE0 e' f
  compareEN' e f@(F y) _ _ = do
    Just f' <- defE y  | Nothing => clashE e f
    compareE0 e f'

  compareEN' e@(B i) f@(B j) _ _ =
    unless (i == j) $ clashE e f
  compareEN' e@(B _) f _ _ = clashE e f

  -- [todo] tracking variance of functions? maybe???
  -- probably not
  compareEN' (fun1 :@ arg1) (fun2 :@ arg2) _ _ =
    local {mode := Equal} $ do
      compareE0 fun1 fun2
      compareT0 arg1 arg2
  compareEN' e@(_ :@ _) f _ _ = clashE e f

  -- [todo] is always checking the types are equal correct?
  compareEN' (tm1 :# ty1) (tm2 :# ty2) _ _ = do
    compareT0 tm1 tm2
    local {mode := Equal} $ compareT0 ty1 ty2
  compareEN' e@(_ :# _) f _ _ = clashE e f

  compareEN' (CloE  {}) _ pe _ = void $ pe IsCloE
  compareEN' (DCloE {}) _ pe _ = void $ pe IsDCloE


  private covering %inline
  compareTN : CanEqual' q _ m => Eq q =>
              NonRedexTerm q 0 n -> NonRedexTerm q 0 n -> m ()
  compareTN s t = compareTN' s.fst t.fst s.snd t.snd

  private covering %inline
  compareEN : CanEqual' q _ m => Eq q =>
              NonRedexElim q 0 n -> NonRedexElim q 0 n -> m ()
  compareEN e f = compareEN' e.fst f.fst e.snd f.snd


  export covering %inline
  compareT : CanEqual' q _ m => Eq q =>
             DimEq d -> Term q d n -> Term q d n -> m ()
  compareT eqs s t =
    for_ (splits eqs) $ \th => compareT0 (s /// th) (t /// th)

  export covering %inline
  compareE : CanEqual' q _ m => Eq q =>
             DimEq d -> Elim q d n -> Elim q d n -> m ()
  compareE eqs e f =
    for_ (splits eqs) $ \th => compareE0 (e /// th) (f /// th)


  export covering %inline
  compareT0 : CanEqual' q _ m => Eq q => Term q 0 n -> Term q 0 n -> m ()
  compareT0 s t = compareTN (whnfT s) (whnfT t)

  export covering %inline
  compareE0 : CanEqual' q _ m => Eq q => Elim q 0 n -> Elim q 0 n -> m ()
  compareE0 e f = compareEN (whnfE e) (whnfE f)

  export covering %inline
  compareTS0 : CanEqual' q _ m => Eq q =>
               ScopeTerm q 0 n -> ScopeTerm q 0 n -> m ()
  compareTS0 (TUnused body1) (TUnused body2) =
    compareT0 body1 body2
  compareTS0 body1 body2 =
    compareT0 (fromScopeTerm body1) (fromScopeTerm body2)


private %inline
into : HasErr q m => HasDefs' q isg m => Eq q =>
       (forall n. HasErr q n => HasEnv' q isg n => d -> a -> a -> n ()) ->
       EqMode -> d -> a -> a -> m ()
into f mode eqs a b =
  runReaderT {m} (MakeEnv {mode, defs = !ask}) $ f eqs a b

export covering %inline
equalTWith : HasErr q m => HasDefs' q _ m => Eq q =>
             DimEq d -> Term q d n -> Term q d n -> m ()
equalTWith = into compareT Equal

export covering %inline
equalEWith : HasErr q m => HasDefs' q _ m => Eq q =>
             DimEq d -> Elim q d n -> Elim q d n -> m ()
equalEWith = into compareE Equal

export covering %inline
subTWith : HasErr q m => HasDefs' q _ m => Eq q =>
           DimEq d -> Term q d n -> Term q d n -> m ()
subTWith = into compareT Sub

export covering %inline
subEWith : HasErr q m => HasDefs' q _ m => Eq q =>
           DimEq d -> Elim q d n -> Elim q d n -> m ()
subEWith = into compareE Sub


export covering %inline
equalT : HasErr q m => HasDefs' q _ m => Eq q =>
         {d : Nat} -> Term q d n -> Term q d n -> m ()
equalT = equalTWith DimEq.new

export covering %inline
equalE : HasErr q m => HasDefs' q _ m => Eq q =>
         {d : Nat} -> Elim q d n -> Elim q d n -> m ()
equalE = equalEWith DimEq.new

export covering %inline
subT : HasErr q m => HasDefs' q _ m => Eq q =>
       {d : Nat} -> Term q d n -> Term q d n -> m ()
subT = subTWith DimEq.new

export covering %inline
subE : HasErr q m => HasDefs' q _ m => Eq q =>
       {d : Nat} -> Elim q d n -> Elim q d n -> m ()
subE = subEWith DimEq.new