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)
| ExpectedW    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
  (ExpectedW       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


||| separates out all the error context layers
||| (e.g. "while checking s : A, …")
export
explodeContext : Error -> (List (Error -> Error), Error)
explodeContext (WhileChecking ctx x s t err) =
  mapFst (WhileChecking ctx x s t ::) $ explodeContext err
explodeContext (WhileCheckingTy ctx s k err) =
  mapFst (WhileCheckingTy ctx s k ::) $ explodeContext err
explodeContext (WhileInferring ctx x e err) =
  mapFst (WhileInferring ctx x e ::) $ explodeContext err
explodeContext (WhileComparingT ctx x s t r err) =
  mapFst (WhileComparingT ctx x s t r ::) $ explodeContext err
explodeContext (WhileComparingE ctx x e f err) =
  mapFst (WhileComparingE ctx x e f ::) $ explodeContext err
explodeContext err = ([], 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 -> String
prettyMode Equal = "equal to"
prettyMode Sub   = "a subtype of"
prettyMode Super = "a supertype of"

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

private
isTypeInUniverse : Maybe Universe -> String
isTypeInUniverse Nothing  = "is a type"
isTypeInUniverse (Just k) = "is a type in universe \{show k}"


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 : {opts : _} -> TyContext d n ->
                BContext n' -> List (QOutput n', Term d n) ->
                Eff Pretty (List (Doc opts))
printCaseQtys ctx ns qts =
  let Evidence _ (ns, qts) = filterSameQtys ns qts in
  traverse (line ns) qts
where
  line : BContext l -> (QOutput l, Term d n) -> Eff Pretty (Doc opts)
  line ns (qs, t) = map (("-" <++>) . sep) $ sequence
    [hangDSingle "the term"
      !(prettyTerm ctx.dnames ctx.tnames t),
     hangDSingle "uses variables" $
      separateTight !commaD $ toSnocList' !(traverse prettyTBind ns),
     hangDSingle "with quantities" $
      separateTight !commaD $ toSnocList' !(traverse prettyQty qs)]

export
prettyErrorNoLoc : {opts : _} -> (showContext : Bool) -> Error ->
                   Eff Pretty (Doc opts)
prettyErrorNoLoc showContext = \case
  ExpectedTYPE _ ctx s =>
    hangDSingle "expected a type universe, but got"
      !(prettyTerm ctx.dnames ctx.tnames s)

  ExpectedPi _ ctx s =>
    hangDSingle "expected a function type, but got"
      !(prettyTerm ctx.dnames ctx.tnames s)

  ExpectedSig _ ctx s =>
    hangDSingle "expected a pair type, but got"
      !(prettyTerm ctx.dnames ctx.tnames s)

  ExpectedW _ ctx s =>
    hangDSingle "expected an inductive (W) type, but got"
      !(prettyTerm ctx.dnames ctx.tnames s)

  ExpectedEnum _ ctx s =>
    hangDSingle "expected an enumeration type, but got"
      !(prettyTerm ctx.dnames ctx.tnames s)

  ExpectedEq _ ctx s =>
    hangDSingle "expected an enumeration type, but got"
      !(prettyTerm ctx.dnames ctx.tnames s)

  ExpectedNat _ ctx s =>
    hangDSingle
      ("expected the type" <++>
        !(prettyTerm [<] [<] $ Nat noLoc) <+> ", but got")
      !(prettyTerm ctx.dnames ctx.tnames s)

  ExpectedBOX _ ctx s =>
    hangDSingle "expected a box type, but got"
      !(prettyTerm ctx.dnames ctx.tnames s)

  BadUniverse _ k l => pure $
    sep ["the universe level" <++> !(prettyUniverse k),
         "is not strictly less than" <++> !(prettyUniverse l)]

  TagNotIn _ tag set =>
    hangDSingle (hsep ["the tag", !(prettyTag tag), "is not contained in"])
      !(prettyTerm [<] [<] $ Enum set noLoc)

  BadCaseEnum _ head body => sep <$> sequence
    [hangDSingle "case expression has head of type"
       !(prettyTerm [<] [<] $ Enum head noLoc),
     hangDSingle "but cases for"
       !(prettyTerm [<] [<] $ Enum body noLoc)]

  BadQtys _ what ctx arms =>
    hangDSingle (text "inconsistent variable usage in \{what}") $
      sep !(printCaseQtys ctx ctx.tnames arms)

  ClashT _ ctx mode ty s t =>
    inEContext ctx . sep =<< sequence
    [hangDSingle "the term" !(prettyTerm [<] ctx.tnames s),
     hangDSingle (text "is not \{prettyMode mode}")
        !(prettyTerm [<] ctx.tnames t),
     hangDSingle "at type" !(prettyTerm [<] ctx.tnames ty)]

  ClashTy _ ctx mode a b =>
    inEContext ctx . sep =<< sequence
    [hangDSingle "the type" !(prettyTerm [<] ctx.tnames a),
     hangDSingle (text "is not \{prettyMode mode}")
        !(prettyTerm [<] ctx.tnames b)]

  ClashE _ ctx mode e f =>
    inEContext ctx . sep =<< sequence
    [hangDSingle "the term" !(prettyElim [<] ctx.tnames e),
     hangDSingle (text "is not \{prettyMode mode}")
        !(prettyElim [<] ctx.tnames f)]

  ClashU _ mode k l => pure $
    sep ["the universe level" <++> !(prettyUniverse k),
         text "is not \{prettyModeU mode}" <++> !(prettyUniverse l)]

  ClashQ _ pi rh => pure $
    sep ["the quantity" <++> !(prettyQty pi),
         "is not equal to" <++> !(prettyQty rh)]

  NotInScope _ x => pure $
    hsep [!(prettyFree x), "is not in scope"]

  NotType _ ctx s =>
    inTContext ctx . sep =<< sequence
    [hangDSingle "the term" !(prettyTerm ctx.dnames ctx.tnames s),
     pure "is not a type"]

  WrongType _ ctx ty s =>
    inEContext ctx . sep =<< sequence
    [hangDSingle "the term" !(prettyTerm [<] ctx.tnames s),
     hangDSingle "cannot have type" !(prettyTerm [<] ctx.tnames ty)]

  MissingEnumArm _ tag tags => pure $
    sep [hsep ["the tag", !(prettyTag tag), "is not contained in"],
         !(prettyTerm [<] [<] $ Enum (fromList tags) noLoc)]

  WhileChecking ctx pi s a err =>
    [|vappendBlank
        (inTContext ctx . sep =<< sequence
          [hangDSingle "while checking" !(prettyTerm ctx.dnames ctx.tnames s),
           hangDSingle "has type"       !(prettyTerm ctx.dnames ctx.tnames a),
           hangDSingle "with quantity"  !(prettyQty pi)])
        (prettyErrorNoLoc showContext err)|]

  WhileCheckingTy ctx a k err =>
    [|vappendBlank
        (inTContext ctx . sep =<< sequence
          [hangDSingle "while checking" !(prettyTerm ctx.dnames ctx.tnames a),
           pure $ text $ isTypeInUniverse k])
        (prettyErrorNoLoc showContext err)|]

  WhileInferring ctx pi e err =>
    [|vappendBlank
        (inTContext ctx . sep =<< sequence
          [hangDSingle "while inferring the type of"
              !(prettyElim ctx.dnames ctx.tnames e),
           hangDSingle "with quantity"  !(prettyQty pi)])
        (prettyErrorNoLoc showContext err)|]

  WhileComparingT ctx mode a s t err =>
    [|vappendBlank
        (inEContext ctx . sep =<< sequence
          [hangDSingle "while checking that" !(prettyTerm [<] ctx.tnames s),
           hangDSingle (text "is \{prettyMode mode}")
              !(prettyTerm [<] ctx.tnames t),
           hangDSingle "at type" !(prettyTerm [<] ctx.tnames a)])
        (prettyErrorNoLoc showContext err)|]

  WhileComparingE ctx mode e f err =>
    [|vappendBlank
        (inEContext ctx . sep =<< sequence
          [hangDSingle "while checking that" !(prettyElim [<] ctx.tnames e),
           hangDSingle (text "is \{prettyMode mode}")
              !(prettyElim [<] ctx.tnames f)])
        (prettyErrorNoLoc showContext err)|]

where
  vappendBlank : Doc opts -> Doc opts -> Doc opts
  vappendBlank a b = flush a `vappend` b

  inTContext : TyContext d n -> Doc opts -> Eff Pretty (Doc opts)
  inTContext ctx doc =
    if showContext && not (null ctx) then
      pure $ vappend doc (sep ["in context", !(prettyTyContext ctx)])
    else pure doc

  inEContext : EqContext n -> Doc opts -> Eff Pretty (Doc opts)
  inEContext ctx doc =
    if showContext && not (null ctx) then
      pure $ vappend doc (sep ["in context", !(prettyEqContext ctx)])
    else pure doc

export
prettyError : {opts : _} -> (showContext : Bool) ->
              Error -> Eff Pretty (Doc opts)
prettyError showContext err = sep <$> sequence
  [prettyLoc err.loc, indentD =<< prettyErrorNoLoc showContext err]