module Quox.Typing.Error

import Quox.Syntax
import Quox.Reduce
import Quox.Typing.Context
import Quox.Typing.EqMode
import Quox.Pretty

import Data.List
import public Control.Monad.Either


public export
data Error q
= ExpectedTYPE (TyContext q d n) (Term q d n)
| ExpectedPi   (TyContext q d n) (Term q d n)
| ExpectedSig  (TyContext q d n) (Term q d n)
| ExpectedEnum (TyContext q d n) (Term q d n)
| ExpectedEq   (TyContext q d n) (Term q d n)
| ExpectedNat  (TyContext q d n) (Term q d n)
| BadUniverse  Universe Universe
| TagNotIn     TagVal (SortedSet TagVal)
| BadCaseEnum  (SortedSet TagVal) (SortedSet TagVal)
| BadCaseQtys  (TyContext q d n) (List (QOutput q n, Term q d n))

-- first term arg of ClashT is the type
| ClashT  (EqContext q n) EqMode (Term q 0 n) (Term q 0 n) (Term q 0 n)
| ClashTy (EqContext q n) EqMode (Term q 0 n) (Term q 0 n)
| ClashE  (EqContext q n) EqMode (Elim q 0 n) (Elim q 0 n)
| ClashU  EqMode Universe Universe
| ClashQ  q q
| NotInScope Name

| NotType (TyContext q d n) (Term q d n)
| WrongType (EqContext q n) (Term q 0 n) (Term q 0 n)

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

| WhnfError WhnfError
%name Error err

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


||| whether the error is surrounded in some context
||| (e.g. "while checking s : A, …")
public export
isErrorContext : Error q -> Bool
isErrorContext (WhileChecking   {}) = True
isErrorContext (WhileCheckingTy {}) = True
isErrorContext (WhileInferring  {}) = True
isErrorContext (WhileComparingT {}) = True
isErrorContext (WhileComparingE {}) = True
isErrorContext _                    = False

||| remove one layer of context
export
peelContext : (e : Error q) -> (0 _ : So (isErrorContext e)) =>
              (Error q -> Error q, Error q)
peelContext (WhileChecking ctx x s t err) =
  (WhileChecking ctx x s t, err)
peelContext (WhileCheckingTy ctx s k err) =
  (WhileCheckingTy ctx s k, err)
peelContext (WhileInferring ctx x e err) =
  (WhileInferring ctx x e, err)
peelContext (WhileComparingT ctx x s t r err) =
  (WhileComparingT ctx x s t r, err)
peelContext (WhileComparingE ctx x e f err) =
  (WhileComparingE ctx x e f, err)

||| separates out all the error context layers
||| (e.g. "while checking s : A, …")
export
explodeContext : Error q -> (List (Error q -> Error q), Error q)
explodeContext err =
  case choose $ isErrorContext err of
    Left y =>
      let (f,  inner) = peelContext err
          (fs, root)  = explodeContext $ assert_smaller err inner in
      (f :: fs, root)
    Right n => ([], err)

||| leaves the outermost context layer, and the innermost (up to) n, and removes
||| the rest if there are more than n+1 in total
export
trimContext : Nat -> Error q -> Error q
trimContext n err =
  case explodeContext err of
    ([],      err) => err
    (f :: fs, err) => f $ foldr apply err $ takeEnd n fs
where
  takeEnd : Nat -> List a -> List a
  takeEnd n = reverse . take n . reverse


export %inline
wrapErr : MonadError e m => (e -> e) -> m a -> m a
wrapErr f act = catchError act $ throwError . f

expect : MonadError e m => (a -> a -> e) -> (a -> a -> Bool) -> a -> a -> m ()
expect err cmp x y = unless (x `cmp` y) $ throwError $ err x y

parameters {auto _ : HasErr q m}
  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


private
prettyMode : EqMode -> Doc HL
prettyMode Equal = "equal to"
prettyMode Sub   = "a subtype of"
prettyMode Super = "a supertype of"

private
prettyModeU : EqMode -> Doc HL
prettyModeU Equal = "equal to"
prettyModeU Sub   = "less than or equal to"
prettyModeU Super = "greater than or equal to"

private
isTypeInUniverse : Maybe Universe -> Doc HL
isTypeInUniverse Nothing  = "is a type"
isTypeInUniverse (Just k) = "is a type in universe" <++> prettyUniverse k

parameters {auto _ : (Eq q, IsQty q, PrettyHL q)} (unicode : Bool)
  private
  termt : TyContext q d n -> Term q d n -> Doc HL
  termt ctx = hang 4 . prettyTerm unicode ctx.dnames ctx.tnames

  private
  terme : EqContext q n -> Term q 0 n -> Doc HL
  terme ctx = hang 4 . prettyTerm unicode [<] ctx.tnames

  private
  dissectCaseQtys : TyContext q d n ->
                    NContext n' -> List (QOutput q n', z) ->
                    List (Doc HL)
  dissectCaseQtys ctx [<]        arms = []
  dissectCaseQtys ctx (tel :< x) arms =
    let qs = map (head . fst) arms
        tl = dissectCaseQtys ctx tel (map (mapFst tail) arms)
    in
    if allSame qs then tl else
    ("-" <++> asep [hsep [pretty0 unicode $ TV x, "is used with quantities"],
                    hseparate comma $ map (pretty0 unicode) qs]) :: tl
  where
    allSame : List q -> Bool
    allSame []        = True
    allSame (q :: qs) = all (== q) qs

  export
  prettyWhnfError : WhnfError -> Doc HL
  prettyWhnfError = \case
    MissingEnumArm tag tags =>
      sep [hsep ["the tag", hl Tag $ pretty tag, "is not contained in"],
           termt empty $ Enum $ fromList tags]

  -- [todo] only show some contexts, probably
  export
  prettyError : (showContext : Bool) -> Error q -> Doc HL
  prettyError showContext = \case
    ExpectedTYPE ctx s =>
      sep ["expected a type universe, but got", termt ctx s]

    ExpectedPi ctx s =>
      sep ["expected a function type, but got", termt ctx s]

    ExpectedSig ctx s =>
      sep ["expected a pair type, but got", termt ctx s]

    ExpectedEnum ctx s =>
      sep ["expected an enumeration type, but got", termt ctx s]

    ExpectedEq ctx s =>
      sep ["expected an equality type, but got", termt ctx s]

    ExpectedNat ctx s =>
      sep ["expected the type ℕ, but got", termt ctx s]

    BadUniverse k l =>
      sep ["the universe level", prettyUniverse k,
           "is not strictly less than", prettyUniverse l]

    TagNotIn tag set =>
      sep [sep ["tag", prettyTag tag, "is not contained in"],
           termt empty (Enum set)]

    BadCaseEnum type arms =>
      sep ["case expression has head of type", termt empty (Enum type),
           "but cases for", termt empty (Enum arms)]

    BadCaseQtys ctx arms =>
      hang 4 $ sep $
        "inconsistent variable usage in case arms" ::
        dissectCaseQtys ctx ctx.tnames arms

    ClashT ctx mode ty s t =>
      inEContext ctx $
      sep ["the term", terme ctx s,
           hsep ["is not", prettyMode mode], terme ctx t,
           "at type", terme ctx ty]

    ClashTy ctx mode a b =>
      inEContext ctx $
      sep ["the type", terme ctx a,
           hsep ["is not", prettyMode mode], terme ctx b]

    ClashE ctx mode e f =>
      inEContext ctx $
      sep ["the term", terme ctx $ E e,
           hsep ["is not", prettyMode mode], terme ctx $ E f]

    ClashU mode k l =>
      sep ["the universe level", prettyUniverse k,
           hsep ["is not", prettyMode mode], prettyUniverse l]

    ClashQ pi rh =>
      sep ["the quantity", pretty0 unicode pi,
           "is not equal to", pretty0 unicode rh]

    NotInScope x =>
      hsep [hl' Free $ pretty0 unicode x, "is not in scope"]

    NotType ctx s =>
      inTContext ctx $
      sep ["the term", termt ctx s, "is not a type"]

    WrongType ctx ty s =>
      inEContext ctx $
      sep ["the term", terme ctx s,
           "cannot have type", terme ctx ty]

    WhileChecking ctx pi s a err =>
      vsep [inTContext ctx $
              sep ["while checking", termt ctx s,
                   "has type", termt ctx a,
                   hsep ["with quantity", pretty0 unicode pi]],
            prettyError showContext err]

    WhileCheckingTy ctx a k err =>
      vsep [inTContext ctx $
              sep ["while checking", termt ctx a,
                   isTypeInUniverse k],
            prettyError showContext err]

    WhileInferring ctx pi e err =>
      vsep [inTContext ctx $
              sep ["while inferring the type of", termt ctx $ E e,
                   hsep ["with quantity", pretty0 unicode pi]],
            prettyError showContext err]

    WhileComparingT ctx mode a s t err =>
      vsep [inEContext ctx $
              sep ["while checking that", terme ctx s,
                   hsep ["is", prettyMode mode], terme ctx t,
                   "at type", terme ctx a],
            prettyError showContext err]

    WhileComparingE ctx mode e f err =>
      vsep [inEContext ctx $
              sep ["while checking that", terme ctx $ E e,
                   hsep ["is", prettyMode mode], terme ctx $ E f],
            prettyError showContext err]

    WhnfError err => prettyWhnfError err
  where
    inTContext : TyContext q d n -> Doc HL -> Doc HL
    inTContext ctx doc =
      if showContext && not (null ctx) then
        vsep [sep ["in context", prettyTyContext unicode ctx], doc]
      else doc

    inEContext : EqContext q n -> Doc HL -> Doc HL
    inEContext ctx doc =
      if showContext && not (null ctx) then
        vsep [sep ["in context", prettyEqContext unicode ctx], doc]
      else doc