module Quox.Syntax.Term.Pretty

import Quox.Syntax.Term.Base
import Quox.Syntax.Term.Subst
import Quox.Context
import Quox.Pretty

import Data.Vect
import Derive.Prelude

%default total
%language ElabReflection


export
prettyUniverse : {opts : _} -> Universe -> Eff Pretty (Doc opts)
prettyUniverse = hl Universe . text . show


export
prettyTerm : {opts : _} -> BContext d -> BContext n -> Term d n ->
             Eff Pretty (Doc opts)

export
prettyElim : {opts : _} -> BContext d -> BContext n -> Elim d n ->
             Eff Pretty (Doc opts)

private
BTelescope : Nat -> Nat -> Type
BTelescope = Telescope' BindName


private
superscript : String -> String
superscript = pack . map sup . unpack where
  sup : Char -> Char
  sup c = case c of
    '0' => '⁰'; '1' => '¹'; '2' => '²'; '3' => '³'; '4' => '⁴'
    '5' => '⁵'; '6' => '⁶'; '7' => '⁷'; '8' => '⁸'; '9' => '⁹'; _ => c


private
PiBind : Nat -> Nat -> Type
PiBind d n = (Qty, BindName, Term d n)

private
pbname : PiBind d n -> BindName
pbname (_, x, _) = x

private
record SplitPi d n where
  constructor MkSplitPi
  {0 inner : Nat}
  binds : Telescope (PiBind d) n inner
  cod   : Term d inner

private
splitPi : Telescope (PiBind d) n n' -> Term d n' -> SplitPi d n
splitPi binds (Pi qty arg res _) =
  splitPi (binds :< (qty, res.name, arg)) $
    assert_smaller res $ pushSubsts' res.term
splitPi binds cod = MkSplitPi {binds, cod}

private
prettyPiBind1 : {opts : _} ->
                Qty -> BindName -> BContext d -> BContext n -> Term d n ->
                Eff Pretty (Doc opts)
prettyPiBind1 pi (BN Unused _) dnames tnames s =
  hcat <$> sequence
    [prettyQty pi, dotD,
     withPrec Arg $ assert_total prettyTerm dnames tnames s]
prettyPiBind1 pi x dnames tnames s = hcat <$> sequence
  [prettyQty pi, dotD,
   hl Delim $ text "(",
   hsep <$> sequence
     [prettyTBind x, hl Delim $ text ":",
      withPrec Outer $ assert_total prettyTerm dnames tnames s],
   hl Delim $ text ")"]

private
prettyPiBinds : {opts : _} ->
                BContext d -> BContext n ->
                Telescope (PiBind d) n n' ->
                Eff Pretty (SnocList (Doc opts))
prettyPiBinds _ _ [<] = pure [<]
prettyPiBinds dnames tnames (binds :< (q, x, t)) =
  let tnames' = tnames . map pbname binds in
  [|prettyPiBinds dnames tnames binds :<
    prettyPiBind1 q x dnames tnames' t|]


private
SigBind : Nat -> Nat -> Type
SigBind d n = (BindName, Term d n)

private
record SplitSig d n where
  constructor MkSplitSig
  {0 inner : Nat}
  binds : Telescope (SigBind d) n inner
  last  : Term d inner

private
splitSig : Telescope (SigBind d) n n' -> Term d n' -> SplitSig d n
splitSig binds (Sig fst snd _) =
  splitSig (binds :< (snd.name, fst)) $
    assert_smaller snd $ pushSubsts' snd.term
splitSig binds last = MkSplitSig {binds, last}

private
prettySigBind1 : {opts : _} ->
                 BindName -> BContext d -> BContext n -> Term d n ->
                 Eff Pretty (Doc opts)
prettySigBind1 (BN Unused _) dnames tnames s =
  withPrec InTimes $ assert_total prettyTerm dnames tnames s
prettySigBind1 x dnames tnames s = hcat <$> sequence
  [hl Delim $ text "(",
   hsep <$> sequence
     [prettyTBind x, hl Delim $ text ":",
      withPrec Outer $ assert_total prettyTerm dnames tnames s],
   hl Delim $ text ")"]

private
prettySigBinds : {opts : _} ->
                 BContext d -> BContext n ->
                 Telescope (SigBind d) n n' ->
                 Eff Pretty (SnocList (Doc opts))
prettySigBinds _ _ [<] = pure [<]
prettySigBinds dnames tnames (binds :< (x, t)) =
  let tnames' = tnames . map fst binds in
  [|prettySigBinds dnames tnames binds :<
    prettySigBind1 x dnames tnames' t|]


private
prettyTypeLine : {opts : _} ->
                 BContext d -> BContext n ->
                 DScopeTerm d n ->
                 Eff Pretty (Doc opts)
prettyTypeLine dnames tnames (S _ (N body)) =
  withPrec Arg (prettyTerm dnames tnames body)
prettyTypeLine dnames tnames (S [< i] (Y body)) =
  parens =<< do
    i'  <- prettyDBind i
    ty' <- withPrec Outer $ prettyTerm (dnames :< i) tnames body
    pure $ sep [hsep [i', !darrowD], ty']


private
data NameSort = T | D
%runElab derive "NameSort" [Eq]

private
NameChunks : Type
NameChunks = SnocList (NameSort, SnocList BindName)

private
record SplitLams d n where
  constructor MkSplitLams
  {0 dinner, ninner : Nat}
  dnames : BTelescope d dinner
  tnames : BTelescope n ninner
  chunks : NameChunks
  body   : Term dinner ninner

private
splitLams : Term d n -> SplitLams d n
splitLams s = go [<] [<] [<] (pushSubsts' s)
where
  push : NameSort -> BindName -> NameChunks -> NameChunks
  push s y [<] = [< (s, [< y])]
  push s y (xss :< (s', xs)) =
    if s == s' then xss :< (s', xs :< y)
    else xss :< (s', xs) :< (s, [< y])

  go : BTelescope d d' -> BTelescope n n' ->
       SnocList (NameSort, SnocList BindName) ->
       Term d' n' -> SplitLams d n
  go dnames tnames chunks (Lam b _) =
    go dnames (tnames :< b.name) (push T b.name chunks) $
       assert_smaller b $ pushSubsts' b.term
  go dnames tnames chunks (DLam b _) =
    go (dnames :< b.name) tnames (push D b.name chunks) $
       assert_smaller b $ pushSubsts' b.term
  go dnames tnames chunks s =
    MkSplitLams dnames tnames chunks s


private
splitTuple : SnocList (Term d n) -> Term d n -> SnocList (Term d n)
splitTuple ss p@(Pair t1 t2 _) =
  splitTuple (ss :< t1) $ assert_smaller p $ pushSubsts' t2
splitTuple ss t = ss :< t


private
prettyTArg : {opts : _} -> BContext d -> BContext n ->
             Term d n -> Eff Pretty (Doc opts)
prettyTArg dnames tnames s =
  withPrec Arg $ assert_total prettyTerm dnames tnames s

private
prettyDArg : {opts : _} -> BContext d -> Dim d -> Eff Pretty (Doc opts)
prettyDArg dnames p = [|atD <+> withPrec Arg (prettyDim dnames p)|]

private
splitApps : Elim d n -> (Elim d n, List (Either (Dim d) (Term d n)))
splitApps e = go [] (pushSubsts' e)
where
  go : List (Either (Dim d) (Term d n)) -> Elim d n ->
       (Elim d n, List (Either (Dim d) (Term d n)))
  go xs e@(App  f s _) =
    go (Right s :: xs) $ assert_smaller e $ pushSubsts' f
  go xs e@(DApp f p _) =
    go (Left  p :: xs) $ assert_smaller e $ pushSubsts' f
  go xs s = (s, xs)

private
prettyDTApps : {opts : _} ->
               BContext d -> BContext n ->
               Elim d n -> List (Either (Dim d) (Term d n)) ->
               Eff Pretty (Doc opts)
prettyDTApps dnames tnames f xs = do
  f  <- withPrec Arg $ assert_total prettyElim dnames tnames f
  xs <- for xs $ either (prettyDArg dnames) (prettyTArg dnames tnames)
  parensIfM App =<< prettyAppD f xs


private
record CaseArm opts d n where
  constructor MkCaseArm
  {0 dinner, ninner : Nat}
  pat    : Doc opts
  dbinds : BTelescope d dinner -- 🍴
  tbinds : BTelescope n ninner
  body   : Term dinner ninner

parameters {opts : LayoutOpts} (dnames : BContext d) (tnames : BContext n)
  private
  prettyCaseArm : CaseArm opts d n -> Eff Pretty (Doc opts)
  prettyCaseArm (MkCaseArm pat dbinds tbinds body) = do
    body <- withPrec Outer $ assert_total
            prettyTerm (dnames . dbinds) (tnames . tbinds) body
    header <- (pat <++>) <$> darrowD
    pure $ ifMultiline (header <++> body) (vsep [header, !(indentD body)])

  private
  prettyCaseBody : List (CaseArm opts d n) -> Eff Pretty (List (Doc opts))
  prettyCaseBody xs = traverse prettyCaseArm xs

private
prettyCompPat : {opts : _} -> DimConst -> BindName -> Eff Pretty (Doc opts)
prettyCompPat e x = [|prettyDimConst e <++> prettyDBind x|]

private
prettyCompArm : {opts : _} -> BContext d -> BContext n ->
                DimConst -> DScopeTerm d n -> Eff Pretty (Doc opts)
prettyCompArm dnames tnames e s = prettyCaseArm dnames tnames $
  MkCaseArm !(prettyCompPat e s.name) [< s.name] [<] s.term

private
layoutComp : {opts : _} ->
             (typq : List (Doc opts)) -> (val, r : Doc opts) ->
             (arms : List (Doc opts)) -> Eff Pretty (Doc opts)
layoutComp typq val r arms = do
  comp <- compD; lb <- hl Delim "{"; rb <- hl Delim "}"
  ind  <- askAt INDENT
  pure $ ifMultiline
    (hsep $ concat {t = List} [[comp], typq, [val, r, lb], arms, [rb]]) $
    (comp <++>
      vsep [sep typq, val, r <++> lb, indent ind $ vsep arms, rb]) <|>
    (vsep $ (comp ::) $ map (indent ind) $ concat {t = List}
      [typq, [val, r <++> lb], map (indent ind) arms, [rb]])


export
prettyEnum : {opts : _} -> List String -> Eff Pretty (Doc opts)
prettyEnum cases =
  tightBraces =<<
  fillSeparateTight !commaD <$>
  traverse (hl Constant . Doc.text . quoteTag) cases

private
prettyCaseRet : {opts : _} ->
                BContext d -> BContext n ->
                ScopeTerm d n -> Eff Pretty (Doc opts)
prettyCaseRet dnames tnames body = withPrec Outer $ case body of
  S _     (N tm) => assert_total prettyTerm dnames tnames tm
  S [< x] (Y tm) => do
    header <- [|prettyTBind x <++> darrowD|]
    body   <- assert_total prettyTerm dnames (tnames :< x) tm
    hangDSingle header body

private
prettyCase_ : {opts : _} ->
              BContext d -> BContext n ->
              Doc opts -> Elim d n -> ScopeTerm d n -> List (CaseArm opts d n) ->
              Eff Pretty (Doc opts)
prettyCase_ dnames tnames intro head ret body = do
  head   <- assert_total prettyElim dnames tnames head
  ret    <- prettyCaseRet dnames tnames ret
  bodys  <- prettyCaseBody dnames tnames body
  return <- returnD; of_ <- ofD
  lb     <- hl Delim "{"; rb <- hl Delim "}"; semi <- semiD
  ind    <- askAt INDENT
  parensIfM Outer $ ifMultiline
    (hsep [intro, head, return, ret, of_, lb, hseparateTight semi bodys, rb])
    (vsep [intro <++> head, return <++> ret, of_ <++> lb,
           indent ind $ vseparateTight semi bodys, rb])

private
prettyCase : {opts : _} ->
             BContext d -> BContext n ->
             Qty -> Elim d n -> ScopeTerm d n -> List (CaseArm opts d n) ->
             Eff Pretty (Doc opts)
prettyCase dnames tnames qty head ret body =
  prettyCase_ dnames tnames ![|caseD <+> prettyQty qty|] head ret body


private
LetBinder : Nat -> Nat -> Type
LetBinder d n = (Qty, BindName, Elim d n)

private
LetExpr : Nat -> Nat -> Nat -> Type
LetExpr d n n' = (Telescope (LetBinder d) n n', Term d n')

private
PrettyLetResult : LayoutOpts -> Nat -> Type
PrettyLetResult opts d =
  Exists $ \n => (BContext n, Term d n, SnocList (Doc opts))

-- [todo] factor out this and the untyped version somehow
export
splitLet : Telescope (LetBinder d) n n' -> Term d n' -> Exists (LetExpr d n)
splitLet ys t@(Let qty rhs body _) =
  splitLet (ys :< (qty, body.name, rhs)) (assert_smaller t body.term)
splitLet ys t =
  Evidence _ (ys, t)

private covering
prettyLets : {opts : LayoutOpts} ->
             BContext d -> BContext a -> Telescope (LetBinder d) a b ->
             Eff Pretty (SnocList (Doc opts))
prettyLets dnames xs lets = sequence $ snd $ go lets where
  go : forall b. Telescope (LetBinder d) a b ->
       (BContext b, SnocList (Eff Pretty (Doc opts)))
  go [<] = (xs, [<])
  go (lets :< (qty, x, rhs)) =
    let (ys, docs) = go lets
        doc = do
          x   <- prettyTBind x
          rhs <- withPrec Outer $ assert_total prettyElim dnames ys rhs
          hangDSingle (hsep [!letD, x, !cstD]) (hsep [rhs, !inD]) in
    (ys :< x, docs :< doc)


private
isDefaultDir : Dim d -> Dim d -> Bool
isDefaultDir (K Zero _) (K One _) = True
isDefaultDir _ _ = False


-- [fixme] use telescopes in Scoped
private
toTel : BContext s -> BTelescope n (s + n)
toTel [<]        = [<]
toTel (ctx :< x) = toTel ctx :< x

private
prettyTyCasePat : {opts : _} ->
                  (k : TyConKind) -> BContext (arity k) ->
                  Eff Pretty (Doc opts)
prettyTyCasePat KTYPE [<] = typeD
prettyTyCasePat KIOState [<] = ioStateD
prettyTyCasePat KPi [< a, b] =
  parens . hsep =<< sequence [prettyTBind a, arrowD, prettyTBind b]
prettyTyCasePat KSig [< a, b] =
  parens . hsep =<< sequence [prettyTBind a, timesD, prettyTBind b]
prettyTyCasePat KEnum [<] = hl Syntax $ text "{}"
prettyTyCasePat KEq [< a0, a1, a, l, r] =
  hsep <$> sequence (eqD :: map prettyTBind [a0, a1, a, l, r])
prettyTyCasePat KNat [<] = natD
prettyTyCasePat KString [<] = stringD
prettyTyCasePat KBOX [< a] = bracks =<< prettyTBind a


prettyLambda : {opts : _} -> BContext d -> BContext n ->
               Term d n -> Eff Pretty (Doc opts)
prettyLambda dnames tnames s =
  parensIfM Outer =<< do
    let MkSplitLams {dnames = ds, tnames = ts, chunks, body} = splitLams s
    hangDSingle !(header chunks)
      !(assert_total prettyTerm (dnames . ds) (tnames . ts) body)
where
  introChar : NameSort -> Eff Pretty (Doc opts)
  introChar T = lamD
  introChar D = dlamD

  prettyBind : NameSort -> BindName -> Eff Pretty (Doc opts)
  prettyBind T = prettyTBind
  prettyBind D = prettyDBind

  header1 : NameSort -> List BindName -> Eff Pretty (Doc opts)
  header1 s xs = hsep <$> sequence
    [introChar s, sep <$> traverse (prettyBind s) xs, darrowD]

  header : NameChunks -> Eff Pretty (Doc opts)
  header cs = sep <$> traverse (\(s, xs) => header1 s (toList xs)) (toList cs)


prettyDisp : {opts : _} -> Universe -> Eff Pretty (Maybe (Doc opts))
prettyDisp 0 = pure Nothing
prettyDisp u = map Just $ hl Universe =<<
  ifUnicode (text $ superscript $ show u) (text $ "^" ++ show u)


prettyTerm dnames tnames (TYPE l _) =
  case !(askAt FLAVOR) of
    Unicode => do
      star  <- hl Syntax "★"
      level <- hl Universe $ text $ superscript $ show l
      pure $ hcat [star, level]
    Ascii   => [|hl Syntax "Type" <++> hl Universe (text $ show l)|]

prettyTerm dnames tnames (IOState _) =
  ioStateD

prettyTerm dnames tnames (Pi qty arg res _) =
  parensIfM Outer =<< do
    let MkSplitPi {binds, cod} = splitPi [< (qty, res.name, arg)] res.term
    arr   <- arrowD
    lines <- map (<++> arr) <$> prettyPiBinds dnames tnames binds
    let tnames = tnames . map pbname binds
    cod <- withPrec Outer $ prettyTerm dnames tnames (assert_smaller res cod)
    pure $ sepSingle $ toList $ lines :< cod

prettyTerm dnames tnames s@(Lam {}) =
  prettyLambda dnames tnames s

prettyTerm dnames tnames (Sig fst snd _) =
  parensIfM Times =<< do
    let MkSplitSig {binds, last} = splitSig [< (snd.name, fst)] snd.term
    times <- timesD
    lines <- map (<++> times) <$> prettySigBinds dnames tnames binds
    let tnames = tnames . map Builtin.fst binds
    last <- withPrec InTimes $
            prettyTerm dnames tnames (assert_smaller snd last)
    pure $ sepSingle $ toList $ lines :< last

prettyTerm dnames tnames p@(Pair fst snd _) =
  parens =<< do
    let elems = splitTuple [< fst] snd
    lines <- for elems $ \t =>
      withPrec Outer $ prettyTerm dnames tnames $ assert_smaller p t
    pure $ separateTight !commaD lines

prettyTerm dnames tnames (Enum cases _) =
  prettyEnum $ SortedSet.toList cases

prettyTerm dnames tnames (Tag tag _) =
  prettyTag tag

prettyTerm dnames tnames (Eq (S _ (N ty)) l r _) =
  parensIfM Eq =<< do
    l  <- withPrec InEq $ prettyTerm dnames tnames l
    r  <- withPrec InEq $ prettyTerm dnames tnames r
    ty <- withPrec InEq $ prettyTerm dnames tnames ty
    pure $ sep [l <++> !eqndD, r <++> !colonD, ty]

prettyTerm dnames tnames (Eq ty l r _) =
  parensIfM App =<< do
    ty <- prettyTypeLine dnames tnames ty
    l  <- withPrec Arg $ prettyTerm dnames tnames l
    r  <- withPrec Arg $ prettyTerm dnames tnames r
    prettyAppD !eqD [ty, l, r]

prettyTerm dnames tnames s@(DLam {}) =
  prettyLambda dnames tnames s

prettyTerm dnames tnames (NAT _) = natD
prettyTerm dnames tnames (Nat n _) = hl Syntax $ pshow n
prettyTerm dnames tnames (Succ p _) =
  parensIfM App =<<
  prettyAppD !succD [!(withPrec Arg $ prettyTerm dnames tnames p)]

prettyTerm dnames tnames (STRING _) = stringD
prettyTerm dnames tnames (Str s _) = prettyStrLit s

prettyTerm dnames tnames (BOX qty ty _) =
  bracks . hcat =<<
    sequence [prettyQty qty, dotD,
              withPrec Outer $ prettyTerm dnames tnames ty]

prettyTerm dnames tnames (Box val _) =
  bracks =<< withPrec Outer (prettyTerm dnames tnames val)

prettyTerm dnames tnames (Let qty rhs body _) = do
  let Evidence _ (lets, body) = splitLet [< (qty, body.name, rhs)] body.term
  heads <- prettyLets dnames tnames lets
  let tnames = tnames . map (\(_, x, _) => x) lets
  body <- withPrec Outer $ assert_total prettyTerm dnames tnames body
  let lines = toList $ heads :< body
  pure $ ifMultiline (hsep lines) (vsep lines)

prettyTerm dnames tnames (E e) = prettyElim dnames tnames e

prettyTerm dnames tnames t0@(CloT (Sub t ph)) =
  prettyTerm dnames tnames $ assert_smaller t0 $ pushSubstsWith' id ph t

prettyTerm dnames tnames t0@(DCloT (Sub t ph)) =
  prettyTerm dnames tnames $ assert_smaller t0 $ pushSubstsWith' ph id t

prettyElim dnames tnames (F x u _) = do
  x <- prettyFree x
  u <- prettyDisp u
  pure $ maybe x (x <+>) u

prettyElim dnames tnames (B i _) =
  prettyTBind $ tnames !!! i

prettyElim dnames tnames e@(App {}) =
  let (f, xs) = splitApps e in
  prettyDTApps dnames tnames f xs

prettyElim dnames tnames (CasePair qty pair ret body _) = do
  let [< x, y] = body.names
  pat <- parens . hsep =<< sequence
         [[|prettyTBind x <+> commaD|], prettyTBind y]
  prettyCase dnames tnames qty pair ret
    [MkCaseArm pat [<] [< x, y] body.term]

prettyElim dnames tnames (Fst pair _) =
  parensIfM App =<< do
    pair <- prettyTArg dnames tnames (E pair)
    prettyAppD !fstD [pair]

prettyElim dnames tnames (Snd pair _) =
  parensIfM App =<< do
    pair <- prettyTArg dnames tnames (E pair)
    prettyAppD !sndD [pair]

prettyElim dnames tnames (CaseEnum qty tag ret arms _) = do
  arms <- for (SortedMap.toList arms) $ \(tag, body) =>
    pure $ MkCaseArm !(prettyTag tag) [<] [<] body
  prettyCase dnames tnames qty tag ret arms

prettyElim dnames tnames (CaseNat qty qtyIH nat ret zero succ _) = do
  let zarm = MkCaseArm !zeroD [<] [<] zero
      [< p, ih] = succ.names
  spat0 <- [|succD <++> prettyTBind p|]
  ihpat0 <- map hcat $ sequence [prettyQty qtyIH, dotD, prettyTBind ih]
  spat <- if ih.val == Unused
          then pure spat0
          else pure $ hsep [spat0 <+> !commaD, ihpat0]
  let sarm = MkCaseArm spat [<] [< p, ih] succ.term
  prettyCase dnames tnames qty nat ret [zarm, sarm]

prettyElim dnames tnames (CaseBox qty box ret body _) = do
  pat <- bracks =<< prettyTBind body.name
  let arm = MkCaseArm pat [<] [< body.name] body.term
  prettyCase dnames tnames qty box ret [arm]

prettyElim dnames tnames e@(DApp {}) =
  let (f, xs) = splitApps e in
  prettyDTApps dnames tnames f xs

prettyElim dnames tnames (Ann tm ty _) =
  parensIfM Outer =<<
    hangDSingle !(withPrec AnnL  [|prettyTerm dnames tnames tm <++> annD|])
                !(withPrec Outer (prettyTerm dnames tnames ty))

prettyElim dnames tnames (Coe ty p q val _) =
  parensIfM App =<< do
    ty  <- prettyTypeLine dnames tnames ty
    p   <- prettyDArg dnames p
    q   <- prettyDArg dnames q
    val <- prettyTArg dnames tnames val
    prettyAppD !coeD [ty, sep [p, q], val]

prettyElim dnames tnames e@(Comp ty p q val r zero one _) =
  parensIfM App =<< do
    ty   <- prettyTypeLine dnames tnames $ assert_smaller e $ SN ty
    pq   <- sep <$> sequence [prettyDArg dnames p, prettyDArg dnames q]
    val  <- prettyTArg dnames tnames val
    r    <- prettyDArg dnames r
    arm0 <- [|prettyCompArm dnames tnames Zero zero <+> semiD|]
    arm1 <- prettyCompArm dnames tnames One one
    ind  <- askAt INDENT
    layoutComp [ty, pq] val r [arm0, arm1]

prettyElim dnames tnames (TypeCase ty ret arms def _) = do
  arms <- for (toList arms) $ \(k ** body) => do
    pat <- prettyTyCasePat k body.names
    pure $ MkCaseArm pat [<] (toTel body.names) body.term
  let darm = MkCaseArm !undD [<] [<] def
  prettyCase_ dnames tnames !typecaseD ty (SN ret) $ arms ++ [darm]

prettyElim dnames tnames e0@(CloE (Sub e ph)) =
  prettyElim dnames tnames $ assert_smaller e0 $ pushSubstsWith' id ph e

prettyElim dnames tnames e0@(DCloE (Sub e ph)) =
  prettyElim dnames tnames $ assert_smaller e0 $ pushSubstsWith' ph id e