module Quox.Typing.Error

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

import Data.List
import Control.Eff


public export
record NameContexts d n where
  constructor MkNameContexts
  dnames : BContext d
  tnames : BContext n

namespace NameContexts
  export
  empty : NameContexts 0 0
  empty = MkNameContexts [<] [<]

  export
  extendDimN : BContext s -> NameContexts d n -> NameContexts (s + d) n
  extendDimN xs = {dnames $= (++ toSnocVect' xs)}

  export
  extendDim : BindName -> NameContexts d n -> NameContexts (S d) n
  extendDim i = extendDimN [< i]

namespace TyContext
  public export
  (.names) : TyContext d n -> NameContexts d n
  (MkTyContext {dnames, tnames, _}).names =
    MkNameContexts {dnames, tnames}

namespace EqContext
  public export
  (.names) : (e : EqContext n) -> NameContexts e.dimLen n
  (MkEqContext {dnames, tnames, _}).names =
    MkNameContexts {dnames, tnames}

  public export
  (.names0) : EqContext n -> NameContexts 0 n
  (MkEqContext {tnames, _}).names0 =
    MkNameContexts {dnames = [<], tnames}

namespace WhnfContext
  public export
  (.names) : WhnfContext d n -> NameContexts d n
  (MkWhnfContext {dnames, tnames, _}).names =
    MkNameContexts {dnames, tnames}


public export
data Error
= ExpectedTYPE Loc (NameContexts d n) (Term d n)
| ExpectedPi   Loc (NameContexts d n) (Term d n)
| ExpectedSig  Loc (NameContexts d n) (Term d n)
| ExpectedEnum Loc (NameContexts d n) (Term d n)
| ExpectedEq   Loc (NameContexts d n) (Term d n)
| ExpectedNat  Loc (NameContexts d n) (Term d n)
| ExpectedBOX  Loc (NameContexts d n) (Term d n)
| BadUniverse  Loc Universe Universe
| TagNotIn     Loc TagVal (SortedSet TagVal)
| BadCaseEnum  Loc (SortedSet TagVal) (SortedSet TagVal)
| BadQtys      Loc String (TyContext d n) (List (QOutput n, Term d n))

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

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

| MissingEnumArm Loc TagVal (List TagVal)

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

public export
ErrorEff : Type -> Type
ErrorEff = Except Error


export
Located Error where
  (ExpectedTYPE    loc _ _).loc       = loc
  (ExpectedPi      loc _ _).loc       = loc
  (ExpectedSig     loc _ _).loc       = loc
  (ExpectedEnum    loc _ _).loc       = loc
  (ExpectedEq      loc _ _).loc       = loc
  (ExpectedNat     loc _ _).loc       = loc
  (ExpectedBOX     loc _ _).loc       = loc
  (BadUniverse     loc _ _).loc       = loc
  (TagNotIn        loc _ _).loc       = loc
  (BadCaseEnum     loc _ _).loc       = loc
  (BadQtys         loc _ _ _).loc     = loc
  (ClashT          loc _ _ _ _ _).loc = loc
  (ClashTy         loc _ _ _ _).loc   = loc
  (ClashE          loc _ _ _ _).loc   = loc
  (ClashU          loc _ _ _).loc     = loc
  (ClashQ          loc _ _).loc       = loc
  (NotInScope      loc _).loc         = loc
  (NotType         loc _ _).loc       = loc
  (WrongType       loc _ _ _).loc     = loc
  (MissingEnumArm  loc _ _).loc       = loc
  (WhileChecking   _ _ _ _ err).loc   = err.loc
  (WhileCheckingTy _ _ _ err).loc     = err.loc
  (WhileInferring  _ _ _ err).loc     = err.loc
  (WhileComparingT _ _ _ _ _ err).loc = err.loc
  (WhileComparingE _ _ _ _ err).loc   = err.loc


||| whether the error is surrounded in some context
||| (e.g. "while checking s : A, …")
public export
isErrorContext : Error -> 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) -> (0 _ : So (isErrorContext e)) =>
              (Error -> Error, Error)
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 -> (List (Error -> Error), Error)
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 -> Error
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


private
expect : Has (Except e) fs =>
         (a -> a -> e) -> (a -> a -> Bool) -> a -> a -> Eff fs ()
expect err cmp x y = unless (x `cmp` y) $ throw $ err x y

parameters {auto _ : Has ErrorEff fs} (loc : Loc)
  export %inline
  expectEqualQ : Qty -> Qty -> Eff fs ()
  expectEqualQ = expect (ClashQ loc) (==)

  export %inline
  expectCompatQ : Qty -> Qty -> Eff fs ()
  expectCompatQ = expect (ClashQ loc) compat

  export %inline
  expectModeU : EqMode -> Universe -> Universe -> Eff fs ()
  expectModeU mode = expect (ClashU loc 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 (unicode : Bool)
  private
  termn : NameContexts d n -> Term d n -> Doc HL
  termn ctx = hang 4 . prettyTerm unicode ctx.dnames ctx.tnames

  private
  dstermn : {s : Nat} -> NameContexts d n -> DScopeTermN s d n -> Doc HL
  dstermn ctx (S i t) = termn (extendDimN i ctx) t.term

  private
  filterSameQtys : BContext n -> List (QOutput n, z) ->
                   Exists $ \n' => (BContext n', List (QOutput n', z))
  filterSameQtys [<]       qts = Evidence 0 ([<], qts)
  filterSameQtys (ns :< n) qts =
    let (qs, qts) = unzip $ map (\(qs :< q, t) => (q, qs, t)) qts
        Evidence l (ns, qts) = filterSameQtys ns qts
    in
    if allSame qs
    then Evidence l (ns, qts)
    else Evidence (S l) (ns :< n, zipWith (\(qs, t), q => (qs :< q, t)) qts qs)
  where
    allSame : List Qty -> Bool
    allSame []        = True
    allSame (q :: qs) = all (== q) qs

  private
  printCaseQtys : TyContext d n ->
                  BContext n' -> List (QOutput n', Term d n) ->
                  List (Doc HL)
  printCaseQtys ctx ns qts =
    let Evidence l (ns, qts) = filterSameQtys ns qts in
    map (line ns) qts
  where
    commaList : PrettyHL a => Context' a l -> Doc HL
    commaList = hseparate comma . map (pretty0 unicode) . toList'

    line : BContext l -> (QOutput l, Term d n) -> Doc HL
    line ns (qs, t) =
      "-" <++> asep ["the term", termn ctx.names t,
                     "uses variables", commaList $ (TV . name) <$> ns,
                     "with quantities", commaList qs]

  export
  prettyErrorNoLoc : (showContext : Bool) -> Error -> Doc HL
  prettyErrorNoLoc showContext = \case
    ExpectedTYPE _ ctx s =>
      sep ["expected a type universe, but got", termn ctx s]

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

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

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

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

    ExpectedNat loc ctx s {d, n} =>
      sep ["expected the type",
           pretty0 unicode $ Nat noLoc {d, n}, "but got", termn ctx s]

    ExpectedBOX loc ctx s =>
      sep ["expected a box type, but got", termn ctx s]

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

    TagNotIn loc tag set =>
      sep [hsep ["tag", prettyTag tag, "is not contained in"],
           termn empty (Enum set noLoc)]

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

    BadQtys loc what ctx arms =>
      hang 4 $ sep $
        hsep ["inconsistent variable usage in", fromString what]
          :: printCaseQtys ctx ctx.tnames arms

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

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

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

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

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

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

    NotType loc ctx s =>
      inTContext ctx $
      sep ["the term", termn ctx.names s, "is not a type"]

    WrongType loc ctx ty s =>
      inEContext ctx $
      sep ["the term", termn ctx.names0 s,
           "cannot have type", termn ctx.names0 ty]

    MissingEnumArm loc tag tags =>
      sep [hsep ["the tag", hl Tag $ pretty tag, "is not contained in"],
           termn empty $ Enum (fromList tags) noLoc]

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

    WhileCheckingTy ctx a k err =>
      vsep [inTContext ctx $
              sep ["while checking", termn ctx.names a,
                   isTypeInUniverse k],
            prettyErrorNoLoc showContext err]

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

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

    WhileComparingE ctx mode e f err =>
      vsep [inEContext ctx $
              sep ["while checking that", termn ctx.names0 $ E e,
                   hsep ["is", prettyMode mode], termn ctx.names0 $ E f],
            prettyErrorNoLoc showContext err]
  where
    inTContext : TyContext 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 n -> Doc HL -> Doc HL
    inEContext ctx doc =
      if showContext && not (null ctx) then
        vsep [sep ["in context", prettyEqContext unicode ctx], doc]
      else doc

  export
  prettyError : (showContext : Bool) -> Error -> Doc HL
  prettyError showContext err =
    sep [prettyLoc err.loc, indent 4 $ prettyErrorNoLoc showContext err]