module Quox.Typing.Error

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

import Data.List
import Control.Eff
import Derive.Prelude

%language ElabReflection
%hide TT.Name

%default total


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

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

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

| WrongBuiltinType Builtin Error
| ExpectedSingleEnum Loc (NameContexts q d n) (Term q d n)

| MissingEnumArm Loc TagVal (List TagVal)

-- extra context
| WhileChecking
    (TyContext q d n) SQty
    (Term q d n) -- term
    (Term q d n) -- type
    Error
| WhileCheckingTy
    (TyContext q d n)
    (Term q d n)
    (Maybe Universe)
    Error
| WhileInferring
    (TyContext q d n) SQty
    (Elim q d n)
    Error
| WhileComparingT
    (EqContext q n) EqMode SQty
    (Term q 0 n) -- type
    (Term q 0 n) (Term q 0 n) -- lhs/rhs
    Error
| WhileComparingE
    (EqContext q n) EqMode SQty
    (Elim q 0 n) (Elim q 0 n)
    Error
%name Error err
-- %runElab derive "Error" [Show]

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
  (ExpectedSTRING   loc _ _).loc        = loc
  (ExpectedBOX      loc _ _).loc        = loc
  (ExpectedIOState  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
  (WrongBuiltinType _ err).loc          = err.loc
  (ExpectedSingleEnum 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 sg s t r err) =
  mapFst (WhileComparingT ctx x sg s t r ::) $ explodeContext err
explodeContext (WhileComparingE ctx x sg e f err) =
  mapFst (WhileComparingE ctx x sg 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) (ctx : BContext q)
  export %inline
  expectEqualQ : Qty q -> Qty q -> Eff fs ()
  expectEqualQ = expect (ClashQ loc ctx) (==)
    -- [fixme] probably replace (==)

  export %inline
  expectCompatQ : Qty q -> Qty q -> Eff fs ()
  expectCompatQ = expect (ClashQ loc ctx) 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 q n, z) ->
                 Exists $ \n' => (BContext n', List (QOutput q 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 q) -> Bool
  allSame []        = True
  allSame (q :: qs) = all (== q) qs


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

parameters {opts : LayoutOpts} (showContext : Bool)
  export
  inContext' : Bool -> a -> (a -> Eff Pretty (Doc opts)) ->
               Doc opts -> Eff Pretty (Doc opts)
  inContext' null ctx f doc =
    if showContext && not null then
      vappend doc <$> hangDSingle "in context" !(f ctx)
    else pure doc

  export %inline
  inTContext : TyContext q d n -> Doc opts -> Eff Pretty (Doc opts)
  inTContext ctx = inContext' (null ctx) ctx prettyTyContext

  export %inline
  inEContext : EqContext q n -> Doc opts -> Eff Pretty (Doc opts)
  inEContext ctx = inContext' (null ctx) ctx prettyEqContext

  export
  prettyErrorNoLoc : Error -> Eff Pretty (Doc opts)
  prettyErrorNoLoc err0 = case err0 of
    ExpectedTYPE _ ctx s =>
      hangDSingle "expected a type universe, but got"
        !(prettyTerm ctx s)

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

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

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

    ExpectedEq _ ctx s =>
      hangDSingle "expected an equality type, but got"
        !(prettyTerm ctx s)

    ExpectedNAT _ ctx s =>
      hangDSingle
        ("expected the type" <++>
          !(prettyTerm ctx $ NAT noLoc) <+> ", but got")
        !(prettyTerm ctx s)

    ExpectedSTRING _ ctx s =>
      hangDSingle
        ("expected the type" <++>
          !(prettyTerm ctx $ STRING noLoc) <+> ", but got")
        !(prettyTerm ctx s)

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

    ExpectedIOState _ ctx s =>
      hangDSingle "expected IOState, but got"
        !(prettyTerm ctx 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 empty $ Enum set noLoc)

    BadCaseEnum _ head body => sep <$> sequence
      [hangDSingle "case expression has head of type"
         !(prettyTerm empty $ Enum head noLoc),
       hangDSingle "but cases for"
         !(prettyTerm empty $ 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 =>
      let names = ctx.names0 in
      inEContext ctx . sep =<< sequence
      [hangDSingle "the term" !(prettyTerm names s),
       hangDSingle (text "is not \{prettyMode mode}") !(prettyTerm names t),
       hangDSingle "at type" !(prettyTerm names ty)]

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

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

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

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

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

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

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

    WrongBuiltinType b err => pure $
      vappend
        (sep [sep ["when checking", text $ builtinDesc b],
              sep ["has type", !(builtinTypeDoc b)]])
        !(prettyErrorNoLoc err)

    ExpectedSingleEnum _ ctx s =>
      hangDSingle "expected an enumeration type with one case, but got"
        !(prettyTerm ctx s)

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

    WhileChecking ctx sg s a err =>
      let names = ctx.names in
      [|vappendBlank
          (inTContext ctx . sep =<< sequence
            [hangDSingle "while checking" !(prettyTerm names s),
             hangDSingle "has type"       !(prettyTerm names a),
             hangDSingle "with quantity"  !(prettyQConst sg.qconst)])
          (prettyErrorNoLoc err)|]

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

    WhileInferring ctx sg e err =>
      [|vappendBlank
          (inTContext ctx . sep =<< sequence
            [hangDSingle "while inferring the type of"
                !(prettyElim ctx.names e),
             hangDSingle "with quantity" !(prettyQConst sg.qconst)])
          (prettyErrorNoLoc err)|]

    WhileComparingT ctx mode sg a s t err =>
      let names = ctx.names0 in
      [|vappendBlank
          (inEContext ctx . sep =<< sequence
            [hangDSingle "while checking that" !(prettyTerm names s),
             hangDSingle (text "is \{prettyMode mode}") !(prettyTerm names t),
             hangDSingle "at type" !(prettyTerm names a),
             hangDSingle "with quantity" !(prettyQConst sg.qconst)])
          (prettyErrorNoLoc err)|]

    WhileComparingE ctx mode sg e f err =>
      let names = ctx.names0 in
      [|vappendBlank
          (inEContext ctx . sep =<< sequence
            [hangDSingle "while checking that" !(prettyElim names e),
             hangDSingle (text "is \{prettyMode mode}") !(prettyElim names f),
             hangDSingle "with quantity" !(prettyQConst sg.qconst)])
          (prettyErrorNoLoc err)|]

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

  export
  prettyError : Error -> Eff Pretty (Doc opts)
  prettyError err = hangDSingle
    !(prettyLoc err.loc)
    !(indentD =<< prettyErrorNoLoc err)