module Quox.Typing

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

import public Data.SortedMap
import public Control.Monad.Either
import Generics.Derive

%hide TT.Name
%hide SOP.from
%hide SOP.to
%default total
%language ElabReflection

%default total


public export
TContext : Type -> Nat -> Nat -> Type
TContext q d = Context (Term q d)

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


public export
record TyContext q d n where
  constructor MkTyContext
  dctx : DimEq d
  tctx : TContext q d n

%name TyContext ctx


namespace TContext
  export %inline
  pushD : TContext q d n -> TContext q (S d) n
  pushD tel = map (// shift 1) tel

namespace TyContext
  public export %inline
  empty : {d : Nat} -> TyContext q d 0
  empty = MkTyContext {dctx = new, tctx = [<]}

  export %inline
  extendTyN : Telescope (Term q d) from to ->
              TyContext q d from -> TyContext q d to
  extendTyN ss = {tctx $= (. ss)}

  export %inline
  extendTy : Term q d n -> TyContext q d n -> TyContext q d (S n)
  extendTy s = extendTyN [< s]

  export %inline
  extendDim : TyContext q d n -> TyContext q (S d) n
  extendDim = {dctx $= (:<? Nothing), tctx $= pushD}

  export %inline
  eqDim : Dim d -> Dim d -> TyContext q d n -> TyContext q d n
  eqDim p q = {dctx $= set p q}


namespace QOutput
  parameters {auto _ : IsQty q}
    export %inline
    (+) : QOutput q n -> QOutput q n -> QOutput q n
    (+) = zipWith (+)

    export %inline
    (*) : q -> QOutput q n -> QOutput q n
    (*) pi = map (pi *)

    export %inline
    zeroFor : TyContext q _ n -> QOutput q n
    zeroFor ctx = zero <$ ctx.tctx


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

public export
CheckResult : DimEq d -> Type -> 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


public export
data EqMode = Equal | Sub | Super
%runElab derive "EqMode" [Generic, Meta, Eq, Ord, DecEq, Show]


public export
data Error q
= ExpectedTYPE (Term q d n)
| ExpectedPi   (Term q d n)
| ExpectedSig  (Term q d n)
| ExpectedEq   (Term q d n)
| BadUniverse  Universe Universe

-- first arg of ClashT is the type
| ClashT EqMode (Term q d n) (Term q d n) (Term q d n)
| ClashE EqMode (Elim q d n) (Elim q d n)
| ClashU EqMode Universe     Universe
| ClashQ        q            q
| ClashD        (Dim d)      (Dim d)
| NotInScope    Name

| NotType (Term q d n)
| WrongType (Term q d n) (Term q d n) (Term q d n)

-- extra context
| WhileChecking
    (TyContext q d n) q
    (Term q d n) -- term
    (Term q d n) -- type
    (Error q)
| WhileInferring
    (TyContext q d n) q
    (Elim q d n)
    (Error q)
| WhileComparingT
    (TyContext q d n) EqMode
    (Term q d n) -- type
    (Term q d n) (Term q d n) -- lhs/rhs
    (Error q)
| WhileComparingE
    (TyContext q d n) EqMode
    (Elim q d n) (Elim q d n)
    (Error q)

public export
0 HasErr : Type -> (Type -> Type) -> Type
HasErr q = MonadError (Error q)

export %inline
wrapErr : HasErr q m => (Error q -> Error q) -> m a -> m a
wrapErr f act = catchError act $ throwError . f

export %inline
ucmp : EqMode -> Universe -> Universe -> Bool
ucmp Equal = (==)
ucmp Sub   = (<=)
ucmp Super = (>=)

export %inline
flip : EqMode -> EqMode
flip Equal = Equal
flip Sub   = Super
flip Super = Sub


parameters {auto _ : HasErr q m}
  export %inline
  expect : Eq a => (a -> a -> Error q) -> (a -> a -> Bool) -> a -> a -> m ()
  expect err cmp x y = unless (x `cmp` y) $ throwError $ err x y

  export %inline
  expectEqualQ : Eq q => q -> q -> m ()
  expectEqualQ = expect ClashQ (==)

  export %inline
  expectCompatQ : IsQty q => q -> q -> m ()
  expectCompatQ = expect ClashQ $ \pi, rh => isYes $ pi `compat` rh

  export %inline
  expectModeU : EqMode -> Universe -> Universe -> m ()
  expectModeU mode = expect (ClashU mode) $ ucmp mode

  export %inline
  expectEqualD : Dim d -> Dim d -> m ()
  expectEqualD = expect ClashD (==)


export
lookupFree' : HasErr q m => Definitions' q g -> Name -> m (Definition' q g)
lookupFree' defs x =
  case lookup x defs of
    Just d  => pure d
    Nothing => throwError $ NotInScope x


public export
substCasePairRet : Term q d n -> ScopeTerm q d n -> Term q d (2 + n)
substCasePairRet dty retty =
  let arg = Pair (BVT 0) (BVT 1) :# (dty // fromNat 2) in
  retty.term // (arg ::: shift 2)


parameters {auto _ : HasErr q m} (defs : Definitions' q _)
  export covering %inline
  expectTYPE : Term q d n -> m Universe
  expectTYPE s =
    case fst $ whnf defs s of
      TYPE l => pure l
      _      => throwError $ ExpectedTYPE s

  export covering %inline
  expectPi : Term q d n -> m (q, Term q d n, ScopeTerm q d n)
  expectPi s =
    case fst $ whnf defs s of
      Pi {qty, arg, res, _} => pure (qty, arg, res)
      _                     => throwError $ ExpectedPi s

  export covering %inline
  expectSig : Term q d n -> m (Term q d n, ScopeTerm q d n)
  expectSig s =
    case fst $ whnf defs s of
      Sig {fst, snd, _} => pure (fst, snd)
      _                 => throwError $ ExpectedSig s

  export covering %inline
  expectEq : Term q d n -> m (DScopeTerm q d n, Term q d n, Term q d n)
  expectEq s =
    case fst $ whnf defs s of
      Eq {ty, l, r, _} => pure (ty, l, r)
      _                => throwError $ ExpectedEq s