/- Adapted from lean-training-data -/
import Lean.Elab.InfoTree
import Lean.Elab.Util
import Lean.Parser.Term
import Lean.PrettyPrinter

open Lean

namespace Lean.Elab

private def elaboratorToString : Name → String
  | .anonymous => ""
  | n => s!"⟨{n}⟩ "

/-- The `Syntax` for a `Lean.Elab.Info`, if there is one. -/
protected def Info.stx? : Info → Option Syntax
  | .ofTacticInfo         info => info.stx
  | .ofTermInfo           info => info.stx
  | .ofCommandInfo        info => info.stx
  | .ofMacroExpansionInfo info => info.stx
  | .ofOptionInfo         info => info.stx
  | .ofFieldInfo          info => info.stx
  | .ofCompletionInfo     info => info.stx
  | .ofUserWidgetInfo     info => info.stx
  | .ofCustomInfo         info => info.stx
  | .ofFVarAliasInfo      _    => none
  | .ofFieldRedeclInfo    info => info.stx
  | .ofChoiceInfo         info => info.stx
  | .ofPartialTermInfo    info => info.stx
  | .ofDelabTermInfo      info => info.stx
  | .ofErrorNameInfo      info => info.stx
/-- Is the `Syntax` for this `Lean.Elab.Info` original, or synthetic? -/
protected def Info.isOriginal (i : Info) : Bool :=
  match i.stx? with
  | none => true   -- Somewhat unclear what to do with `FVarAliasInfo`, so be conservative.
  | some stx => match stx.getHeadInfo with
    | .original .. => true
    | _ => false

def ContextInfo.ppExpr (ctx : ContextInfo) (lctx : LocalContext) (e : Expr) : IO Format :=
  ctx.runMetaM lctx (do Meta.ppExpr (← instantiateMVars e))

def CommandInfo.toString (info : CommandInfo) (ctx : ContextInfo) : IO String := do
  let stx := (← ctx.ppSyntax {} info.stx).pretty
  return s!"{elaboratorToString info.elaborator}\n{stx}"

def TermInfo.toString (info : TermInfo) (ctx : ContextInfo) : IO String := do
  let stx := (← ctx.ppSyntax info.lctx info.stx).pretty
  let expectedType := (← info.expectedType?.mapM fun ty => do
    pure s!": {(← ctx.ppExpr info.lctx ty).pretty}").getD ""
  let expr := (← ctx.ppExpr info.lctx info.expr).pretty
  return s!"{elaboratorToString info.elaborator}{expr}{expectedType}\n{stx}"

/-- Find the name for the outermost `Syntax` in this `TacticInfo`. -/
def TacticInfo.name? (t : TacticInfo) : Option Name :=
  match t.stx with
  | Syntax.node _ n _ => some n
  | _ => none
/-- Decide whether a tactic is "substantive",
or is merely a tactic combinator (e.g. `by`, `;`, multiline tactics, parenthesized tactics). -/
def TacticInfo.isSubstantive (t : TacticInfo) : Bool :=
  match t.name? with
  | none => false
  | some `null => false
  | some ``cdot => false
  | some ``cdotTk => false
  | some ``Lean.Parser.Term.byTactic => false
  | some ``Lean.Parser.Tactic.tacticSeq => false
  | some ``Lean.Parser.Tactic.tacticSeq1Indented => false
  | some ``Lean.Parser.Tactic.«tactic_<;>_» => false
  | some ``Lean.Parser.Tactic.paren => false
  | _ => true
def TacticInfo.pp (info : TacticInfo) (ctx : ContextInfo) : IO Format :=
  ctx.runMetaM {} try
    Lean.PrettyPrinter.ppTactic ⟨info.stx⟩
  catch _ =>
    pure "<failed to pretty print>"
def TacticInfo.toString (i : TacticInfo) (ctx : ContextInfo) : IO String := do
  let name := i.name?
  let stx := Format.pretty (← i.pp ctx)
  return s!"{name}\n{stx}"

/--
Keep `.node` nodes and `.hole` nodes satisfying predicates.

Returns a `List InfoTree`, although in most situations this will be a singleton.
-/
partial def InfoTree.filter (p : Info → Bool) (m : MVarId → Bool := fun _ => false) :
    InfoTree → List InfoTree
  | .context ctx tree => tree.filter p m |>.map (.context ctx)
  | .node info children =>
    if p info then
      [.node info (children.toList.map (filter p m)).flatten.toPArray']
    else
      (children.toList.map (filter p m)).flatten
  | .hole mvar => if m mvar then [.hole mvar] else []

/-- Analogue of `Lean.Elab.InfoTree.findInfo?`, but that returns a list of all results. -/
partial def InfoTree.findAllInfo
    (t : InfoTree)
    (context?: Option Elab.ContextInfo)
    (haltOnMatch : Bool := false)
    (pred : Elab.Info → Bool)
    : List (Elab.Info × Option Elab.ContextInfo × PersistentArray Elab.InfoTree) :=
  match t with
  | .context inner t => findAllInfo t (inner.mergeIntoOuter? context?) haltOnMatch pred
  | .node i children  =>
    let head := if pred i then [(i, context?, children)] else []
    let tail := if haltOnMatch ∧ !head.isEmpty then [] else children.toList.flatMap (fun t => findAllInfo t context? haltOnMatch pred)
    head ++ tail
  | _ => []

/-- Monadic analogue of `findAllInfo`, but predicate controls whether to recurse. -/
partial def InfoTree.findAllInfoM [Monad m]
    (t : InfoTree)
    (context?: Option Elab.ContextInfo)
    (pred : Elab.Info → Option Elab.ContextInfo → m (Bool × Bool))
    : m (List (Elab.Info × Option Elab.ContextInfo × PersistentArray Elab.InfoTree)) := do
  match t with
  | .context inner t => t.findAllInfoM (inner.mergeIntoOuter? context?) pred
  | .node i children  =>
    let (flagCollect, flagRecurse) ← pred i context?
    let head := if flagCollect then [(i, context?, children)] else []
    let tail := if ¬ flagRecurse then pure [] else children.toList.mapM (fun t => t.findAllInfoM context? pred)
    return head ++ (← tail).flatten
  | _ => return []

structure InfoTreePrintOptions where
  prettyPrint : Bool := true

@[export pantograph_infotree_to_string_m]
partial def InfoTree.toString (t : InfoTree) (ctx?: Option Elab.ContextInfo := .none) (options : InfoTreePrintOptions := {}) (indent : Nat := 0)
  : IO String := do
  let space := String.join $ List.replicate indent "  "
  match t with
  | .context ctx t =>
    let s ← t.toString (ctx.mergeIntoOuter? ctx?)
    pure s!"[context] {s}"
  | .node info children =>
    if let some ctx := ctx? then
      let node : String ← match info with
      | .ofTermInfo { stx, expr, .. } =>
        pure s!"[term] {stx.prettyPrint} ➡ {expr}"
      | .ofCommandInfo info =>
        let head := match info.stx.getArg 1 with
          | .missing => ""
          | other => s!" {other.getKind.toString}"
        let s := if options.prettyPrint then
            info.stx.prettyPrint
          else
            s!"{info.stx}"
        pure s!"[command/{info.stx.getKind.toString}{head}] {s}"
      | .ofTacticInfo  info => pure s!"[tactic] {info.stx.prettyPrint}"
      | .ofMacroExpansionInfo _ => pure "[macro_exp]"
      | .ofOptionInfo info => pure s!"[option] {info.stx.prettyPrint}"
      | .ofFieldInfo _ => pure "[field]"
      | .ofCompletionInfo info => pure s!"[completion] {info.stx.prettyPrint}"
      | .ofUserWidgetInfo _ => pure "[user_widget]"
      | .ofCustomInfo _ => pure "[custom]"
      | .ofFVarAliasInfo _ => pure "[fvar_alias]"
      | .ofFieldRedeclInfo _ => pure "[field_redecl]"
      | .ofChoiceInfo _ => pure "[choice]"
      | .ofPartialTermInfo  _ => pure "[partial_term]"
      | .ofDelabTermInfo _ => pure "[delab_term]"
      | .ofErrorNameInfo _ => pure "[error_name]"
      let children ← children.toList.mapM λ t' => do
        t'.toString ctx options (indent := indent + 1)
      let children := String.join children
      return s!"{space}{node}\n{children}"
    else throw <| IO.userError "No `ContextInfo` available."
  | .hole mvarId =>
    if let some ctx := ctx? then
      let payload := (← ctx.runMetaM {} (do Meta.ppGoal mvarId)).pretty
      return s!"{space}[hole] {payload}"
    else throw <| IO.userError "No `ContextInfo` available."

end Lean.Elab