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.Whnf
import public Quox.Pretty

import Language.Reflection
import Control.Eff

%default total
%language ElabReflection

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


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

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

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

public export
InferResult : DimEq d -> TermLike
InferResult eqs q d n = IfConsistent eqs $ InferResult' q 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 q d n -> ScopeTerm q d n ->
                   Term q 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 // fromNat 2) tm.loc in
  retty.term // (arg ::: shift 2)

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

public export
substCaseBoxRet : BindName -> Term q d n -> ScopeTerm q d n -> Term q 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)


private
0 ExpectErrorConstructor : Type
ExpectErrorConstructor =
  forall q, d, n. Loc -> NameContexts q d n -> Term q d n -> Error

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

    private covering %macro
    expect : ExpectErrorConstructor -> TTImp -> TTImp ->
             Elab (Term q d n -> Eff fs a)
    expect err pat rhs = Prelude.do
      match <- check `(\case ~(pat) => Just ~(rhs); _ => Nothing)
      pure $ \term => do
        res <- whnf term
        maybe (throw $ err loc ctx.names term) pure $ match $ fst res

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

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

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

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

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

    export covering %inline
    expectNAT : Term q d n -> Eff fs ()
    expectNAT = expect ExpectedNAT `(NAT {}) `(())

    export covering %inline
    expectSTRING : Term q d n -> Eff fs ()
    expectSTRING = expect ExpectedSTRING `(STRING {}) `(())

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

    export covering %inline
    expectIOState : Term q d n -> Eff fs ()
    expectIOState = expect ExpectedIOState `(IOState {}) `(())


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

    private covering %macro
    expect : ExpectErrorConstructor -> TTImp -> TTImp ->
             Elab (Term q 0 n -> Eff fs a)
    expect err pat rhs = do
      match <- check `(\case ~(pat) => Just ~(rhs); _ => Nothing)
      pure $ \term => do
        res <- whnf term
        let t0 = delay $ term // shift0 ctx.dimLen
        maybe (throw $ err loc ctx.names t0) pure $ match $ fst res

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

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

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

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

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

    export covering %inline
    expectNAT : Term q 0 n -> Eff fs ()
    expectNAT = expect ExpectedNAT `(NAT {}) `(())

    export covering %inline
    expectSTRING : Term q 0 n -> Eff fs ()
    expectSTRING = expect ExpectedSTRING `(STRING {}) `(())

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

    export covering %inline
    expectIOState : Term q 0 n -> Eff fs ()
    expectIOState = expect ExpectedIOState `(IOState {}) `(())