module Quox.Typing

import public Quox.Typing.Context as Typing
import public Quox.Typing.EqMode as Typing
import public Quox.Typing.Error as Typing

import public Quox.Syntax
import public Quox.Definition
import public Quox.Reduce

import Language.Reflection
import Control.Eff

%default total
%language ElabReflection

private TTName : Type
TTName = TT.Name
%hide TT.Name


public export
CheckResult' : Nat -> Type
CheckResult' = QOutput

public export
CheckResult : DimEq d -> Nat -> Type
CheckResult eqs n = IfConsistent eqs $ CheckResult' n

public export
record InferResult' d n where
  constructor InfRes
  type : Term d n
  qout : QOutput n

public export
InferResult : DimEq d -> TermLike
InferResult eqs d n = IfConsistent eqs $ InferResult' d n


export
lookupFree : Has ErrorEff fs => Name -> Loc -> Definitions -> Eff fs Definition
lookupFree x loc defs = maybe (throw $ NotInScope loc x) pure $ lookup x defs


public export
substCasePairRet : BContext 2 -> Term d n -> ScopeTerm d n -> Term d (2 + n)
substCasePairRet [< x, y] dty retty =
  let tm = Pair (BVT 1 x.loc) (BVT 0 y.loc) $ x.loc `extendL` y.loc
      arg = Ann tm (dty // shift 2) tm.loc
  in
  retty.term // (arg ::: shift 2)

public export
substCaseWRet : BContext 3 -> Term d n -> ScopeTerm d n -> Term d (3 + n)
substCaseWRet [< x, y, ih] dty retty =
  let tm = Sup (BVT 2 x.loc) (BVT 1 y.loc) $ x.loc `extendL` y.loc
      arg = Ann tm (dty // shift 3) tm.loc
  in
  sub1 (weakS 3 retty) arg

public export
substCaseSuccRet : BContext 2 -> ScopeTerm d n -> Term d (2 + n)
substCaseSuccRet [< p, ih] retty =
  let arg = Ann (Succ (BVT 1 p.loc) p.loc) (Nat noLoc) $ p.loc `extendL` ih.loc
  in
  retty.term // (arg ::: shift 2)

public export
substCaseBoxRet : BindName -> Term d n -> ScopeTerm d n -> Term d (S n)
substCaseBoxRet x dty retty =
  let arg = Ann (Box (BVT 0 x.loc) x.loc) (weakT 1 dty) x.loc in
  retty.term // (arg ::: shift 1)


parameters (defs : Definitions) {auto _ : (Has ErrorEff fs, Has NameGen fs)}
  namespace TyContext
   parameters (ctx : TyContext d n) (loc : Loc)
    export covering
    whnf : {0 isRedex : RedexTest tm} -> CanWhnf tm isRedex =>
           tm d n -> Eff fs (NonRedex tm d n defs)
    whnf tm = do
      let Val n = ctx.termLen; Val d = ctx.dimLen
      res <- lift $ runExcept $ whnf defs (toWhnfContext ctx) tm
      rethrow res

    private covering %macro
    expect : (forall d, n. Loc -> NameContexts d n -> Term d n -> Error) ->
             TTImp -> TTImp -> Elab (Term d n -> Eff fs a)
    expect k l r = do
      f <- check `(\case ~(l) => Just ~(r); _ => Nothing)
      pure $ \t => maybe (throw $ k loc ctx.names t) pure . f . fst =<< whnf t

    export covering %inline
    expectTYPE : Term d n -> Eff fs Universe
    expectTYPE = expect ExpectedTYPE `(TYPE {l, _}) `(l)

    export covering %inline
    expectPi : Term d n -> Eff fs (Qty, Term d n, ScopeTerm d n)
    expectPi = expect ExpectedPi `(Pi {qty, arg, res, _}) `((qty, arg, res))

    export covering %inline
    expectSig : Term d n -> Eff fs (Term d n, ScopeTerm d n)
    expectSig = expect ExpectedSig `(Sig {fst, snd, _}) `((fst, snd))

    export covering %inline
    expectW : Term d n -> Eff fs (Term d n, ScopeTerm d n)
    expectW = expect ExpectedW `(W {shape, body, _}) `((shape, body))

    export covering %inline
    expectEnum : Term d n -> Eff fs (SortedSet TagVal)
    expectEnum = expect ExpectedEnum `(Enum {cases, _}) `(cases)

    export covering %inline
    expectEq : Term d n -> Eff fs (DScopeTerm d n, Term d n, Term d n)
    expectEq = expect ExpectedEq `(Eq {ty, l, r, _}) `((ty, l, r))

    export covering %inline
    expectNat : Term d n -> Eff fs ()
    expectNat = expect ExpectedNat `(Nat {}) `(())

    export covering %inline
    expectBOX : Term d n -> Eff fs (Qty, Term d n)
    expectBOX = expect ExpectedBOX `(BOX {qty, ty, _}) `((qty, ty))


  namespace EqContext
   parameters (ctx : EqContext n) (loc : Loc)
    export covering
    whnf : {0 isRedex : RedexTest tm} -> CanWhnf tm isRedex =>
           tm 0 n -> Eff fs (NonRedex tm 0 n defs)
    whnf tm = do
      let Val n = ctx.termLen
      res <- lift $ runExcept $ whnf defs (toWhnfContext ctx) tm
      rethrow res

    private covering %macro
    expect : (forall d, n. Loc -> NameContexts d n -> Term d n -> Error) ->
             TTImp -> TTImp -> Elab (Term 0 n -> Eff fs a)
    expect k l r = do
      f <- check `(\case ~(l) => Just ~(r); _ => Nothing)
      pure $ \t =>
        let err = throw $ k loc ctx.names (t // shift0 ctx.dimLen) in
        maybe err pure . f . fst =<< whnf t

    export covering %inline
    expectTYPE : Term 0 n -> Eff fs Universe
    expectTYPE = expect ExpectedTYPE `(TYPE {l, _}) `(l)

    export covering %inline
    expectPi : Term 0 n -> Eff fs (Qty, Term 0 n, ScopeTerm 0 n)
    expectPi = expect ExpectedPi `(Pi {qty, arg, res, _}) `((qty, arg, res))

    export covering %inline
    expectSig : Term 0 n -> Eff fs (Term 0 n, ScopeTerm 0 n)
    expectSig = expect ExpectedSig `(Sig {fst, snd, _}) `((fst, snd))

    export covering %inline
    expectW : Term 0 n -> Eff fs (Term 0 n, ScopeTerm 0 n)
    expectW = expect ExpectedW `(W {shape, body, _}) `((shape, body))

    export covering %inline
    expectEnum : Term 0 n -> Eff fs (SortedSet TagVal)
    expectEnum = expect ExpectedEnum `(Enum {cases, _}) `(cases)

    export covering %inline
    expectEq : Term 0 n -> Eff fs (DScopeTerm 0 n, Term 0 n, Term 0 n)
    expectEq = expect ExpectedEq `(Eq {ty, l, r, _}) `((ty, l, r))

    export covering %inline
    expectNat : Term 0 n -> Eff fs ()
    expectNat = expect ExpectedNat `(Nat {}) `(())

    export covering %inline
    expectBOX : Term 0 n -> Eff fs (Qty, Term 0 n)
    expectBOX = expect ExpectedBOX `(BOX {qty, ty, _}) `((qty, ty))