import Mathlib.Init.Data.Nat.Notation
import Mathlib.Lean.Message
import Mathlib.Lean.Expr.Basic
import Mathlib.Data.String.Defs
import Mathlib.Tactic.Simps.NotationClass

open Lean Elab Parser Command
open Meta hiding Config
open Elab.Term hiding mkConst

def updateName (nm : Name) (s : String) (isPrefix : Bool) : Name :=
  nm.updateLast fun s' ↦ if isPrefix then s ++ "_" ++ s' else s' ++ "_" ++ s

namespace Lean.Meta
open Tactic Simp
def mkSimpContextResult (cfg : Meta.Simp.Config := {}) (simpOnly := false) (kind := SimpKind.simp)
    (dischargeWrapper := DischargeWrapper.default) (hasStar := false) :
    MetaM MkSimpContextResult := do
  match dischargeWrapper with
  | .default => pure ()
  | _ =>
    if kind == SimpKind.simpAll then
      throwError "'simp_all' tactic does not support 'discharger' option"
    if kind == SimpKind.dsimp then
      throwError "'dsimp' tactic does not support 'discharger' option"
  let simpTheorems ← if simpOnly then
    simpOnlyBuiltins.foldlM (·.addConst ·) ({} : SimpTheorems)
  else
    getSimpTheorems
  let congrTheorems ← getSimpCongrTheorems
  let ctx : Simp.Context := {
    config      := cfg
    simpTheorems := #[simpTheorems], congrTheorems
  }
  if !hasStar then
    return { ctx, dischargeWrapper }
  else
    let mut simpTheorems := ctx.simpTheorems
    let hs ← getPropHyps
    for h in hs do
      unless simpTheorems.isErased (.fvar h) do
        simpTheorems ← simpTheorems.addTheorem (.fvar h) (← h.getDecl).toExpr
    let ctx := { ctx with simpTheorems }
    return { ctx, dischargeWrapper }

def mkSimpContext (cfg : Meta.Simp.Config := {}) (simpOnly := false) (kind := SimpKind.simp)
    (dischargeWrapper := DischargeWrapper.default) (hasStar := false) :
    MetaM Simp.Context := do
  let data ← mkSimpContextResult cfg simpOnly kind dischargeWrapper hasStar
  return data.ctx

end Lean.Meta

def hasSimpAttribute (env : Environment) (declName : Name) : Bool :=
  simpExtension.getState env |>.lemmaNames.contains <| .decl declName

namespace Lean.Parser
namespace Attr

attribute [notation_class] Add

syntax simpsArgsRest := (Tactic.config)? (ppSpace ident)*

syntax (name := simps) "simps" "?"? simpsArgsRest : attr

macro "simps?" rest:simpsArgsRest : attr => `(attr| simps ? $rest)

end Attr

namespace Command

syntax simpsRule.rename := ident " → " ident
syntax simpsRule.erase := "-" ident
syntax simpsRule.add := "+" ident
syntax simpsRule.prefix := &"as_prefix" ident
syntax simpsRule := simpsRule.prefix <|> simpsRule.rename <|> simpsRule.erase <|> simpsRule.add
syntax simpsProj := (ppSpace ident (" (" simpsRule,+ ")")?)

syntax (name := initialize_simps_projections)
  "initialize_simps_projections" "?"? simpsProj : command

macro "initialize_simps_projections?" rest:simpsProj : command =>
  `(initialize_simps_projections ? $rest)

end Command
end Lean.Parser

initialize registerTraceClass `simps.verbose
initialize registerTraceClass `simps.debug

namespace Simps

structure UnusedStructure where
  name : Name
  expr : Expr
  projNrs : List ℕ
  isDefault : Bool
  isPrefix : Bool
  deriving Inhabited

instance : ToMessageData UnusedStructure where toMessageData
  | ⟨a, b, c, d, e⟩ => .group <| .nest 1 <|
    "⟨" ++ .joinSep [toMessageData a, toMessageData b, toMessageData c, toMessageData d,
      toMessageData e] ("," ++ Format.line) ++ "⟩"

initialize structureExt : NameMapExtension (List Name × Array UnusedStructure) ←
  registerNameMapExtension (List Name × Array UnusedStructure)

structure ParsedProjectionData where
  origName : Name × Syntax
  newName : Name × Syntax
  isDefault : Bool := true
  isPrefix : Bool := false
  expr? : Option Expr := none
  projNrs : Array Nat := #[]
  isCustom : Bool := false

def ParsedProjectionData.toUnusedStructure (p : ParsedProjectionData) : UnusedStructure :=
  { p with name := p.newName.1, expr := p.expr?.getD default, projNrs := p.projNrs.toList }

instance : ToMessageData ParsedProjectionData where toMessageData
  | ⟨x₁, x₂, x₃, x₄, x₅, x₆, x₇⟩ => .group <| .nest 1 <|
    "⟨" ++ .joinSep [toMessageData x₁, toMessageData x₂, toMessageData x₃, toMessageData x₄,
      toMessageData x₅, toMessageData x₆, toMessageData x₇]
    ("," ++ Format.line) ++ "⟩"

inductive ProjectionRule where
  | rename (oldName : Name) (oldStx : Syntax) (newName : Name) (newStx : Syntax) :
      ProjectionRule
  | add : Name → Syntax → ProjectionRule
  | erase : Name → Syntax → ProjectionRule
  | prefix : Name → Syntax → ProjectionRule

instance : ToMessageData ProjectionRule where toMessageData
  | .rename x₁ x₂ x₃ x₄ => .group <| .nest 1 <|
    "rename ⟨" ++ .joinSep [toMessageData x₁, toMessageData x₂, toMessageData x₃, toMessageData x₄]
      ("," ++ Format.line) ++ "⟩"
  | .add x₁ x₂ => .group <| .nest 1 <|
    "+⟨" ++ .joinSep [toMessageData x₁, toMessageData x₂] ("," ++ Format.line) ++ "⟩"
  | .erase x₁ x₂ => .group <| .nest 1 <|
    "-⟨" ++ .joinSep [toMessageData x₁, toMessageData x₂] ("," ++ Format.line) ++ "⟩"
  | .prefix x₁ x₂ => .group <| .nest 1 <|
    "prefix ⟨" ++ .joinSep [toMessageData x₁, toMessageData x₂] ("," ++ Format.line) ++ "⟩"

def projectionsInfo (l : List UnusedStructure) (pref : String) (str : Name) : MessageData :=
  let ⟨defaults, nondefaults⟩ := l.partition (·.isDefault)
  let toPrint : List MessageData :=
    defaults.map fun s ↦
      let prefixStr := if s.isPrefix then "(prefix) " else ""
      m!"Projection {prefixStr}{s.name}: {s.expr}"
  let print2 : MessageData :=
    String.join <| (nondefaults.map fun nm : UnusedStructure ↦ toString nm.1).intersperse ", "
  let toPrint :=
    toPrint ++
      if nondefaults.isEmpty then [] else
      [("No lemmas are generated for the projections: " : MessageData) ++ print2 ++ "."]
  let toPrint := MessageData.joinSep toPrint ("\n" : MessageData)
  m! "{pref} {str}:\n{toPrint}"

def findProjectionIndices (strName projName : Name) : MetaM (List ℕ) := do
  let env ← getEnv
  let .some baseStr := findField? env strName projName |
    throwError "{strName} has no field {projName} in parent structure"
  let .some fullProjName := getProjFnForField? env baseStr projName |
    throwError "no such field {projName}"
  let .some pathToField := getPathToBaseStructure? env baseStr strName |
    throwError "no such field {projName}"
  let allProjs := pathToField ++ [fullProjName]
  return allProjs.map (env.getProjectionFnInfo? · |>.get!.i)

partial def getCompositeOfProjectionsAux (stx : Syntax)
    (proj : String) (e : Expr) (pos : Array ℕ) (args : Array Expr) : MetaM (Expr × Array ℕ) := do
  dbg_trace s!"1"
  let env ← getEnv
  let .const structName _ := (← whnf (←inferType e)).getAppFn |
    throwError "{e} doesn't have a structure as type"
  dbg_trace s!"2 {structName}"
  let projs := getStructureFieldsFlattened env structName
  dbg_trace s!"3 {projs.size}"
  let projInfo := projs.toList.map fun p ↦ do
    (← ("_" ++ p.getString).isPrefixOf? proj, p)
  dbg_trace s!"4 {projInfo.length} {projInfo}"
  let some (projRest, projName) := projInfo.reduceOption.getLast? |
    throwError "Failed to find constructor {proj.drop 1} in structure {structName}."
  dbg_trace "z"
  let newE ← mkProjection e projName
  dbg_trace "4"
  let newPos := pos ++ (← findProjectionIndices structName projName)
  if projRest.isEmpty then
    let newE ← mkLambdaFVars args newE
    if !stx.isMissing then
      _ ← TermElabM.run' <| addTermInfo stx newE
    return (newE, newPos)
  let type ← inferType newE
  forallTelescopeReducing type fun typeArgs _tgt ↦ do
    getCompositeOfProjectionsAux stx projRest (mkAppN newE typeArgs) newPos (args ++ typeArgs)

def mkParsedProjectionData (structName : Name) : CoreM (Array ParsedProjectionData) := do
  let env ← getEnv
  let projs := getStructureFields env structName
  if projs.size == 0 then
    throwError "Declaration {structName} is not a structure."
  let projData := projs.map fun fieldName ↦ {
    origName := (fieldName, .missing), newName := (fieldName, .missing),
    isDefault := isSubobjectField? env structName fieldName |>.isNone }
  let parentProjs := getStructureFieldsFlattened env structName false
  let parentProjs := parentProjs.filter (!projs.contains ·)
  let parentProjData := parentProjs.map fun nm ↦
    {origName := (nm, .missing), newName := (nm, .missing)}
  return projData ++ parentProjData

set_option linter.missingDocs false
/-! the crash seems to only appear -/
/-! if there are exactly two such documentation sections -/

/-- Execute the projection renamings (and turning off projections) as specified by `rules`. -/
def applyProjectionRules (projs : Array ParsedProjectionData) (rules : Array ProjectionRule) :
  CoreM (Array ParsedProjectionData) := do
  let projs : Array ParsedProjectionData := rules.foldl (init := projs) fun projs rule ↦
    match rule with
    | .rename oldName oldStx newName newStx =>
      dbg_trace s!"{oldName} {oldStx} {newName} {newStx}"
      if (projs.map (·.newName.1)).contains oldName then
        projs.map fun proj ↦ if proj.newName.1 == oldName then
          { proj with
            newName := (newName, newStx),
            origName.2 := if proj.origName.2.isMissing then oldStx else proj.origName.2 }
        else
          proj
      else
        projs.push {origName := (oldName, oldStx), newName := (newName, newStx)}
    | _ => projs
  pure projs

#eval liftCoreM <| do
  _ ← applyProjectionRules #[{origName := (`val, .missing), newName := (`val, .missing)}]
    #[.rename `val .missing `coe .missing]