BeginPackage["Undo`"]
SetUpUndo::usage="";
UndoDataNotebook::usage="";

Begin["`Private`"]

UndoDataNotebook[]:=CurrentValue[EvaluationNotebook[], {TaggingRules, "undoNotebook"}]

SetAttributes[withMacros, HoldAllComplete];
withMacros[expr_][nb_][string_List]:=Unevaluated[expr]/.(
string/.s_String:>(s:>CurrentValue[nb, {TaggingRules, "Undo","Settings", s}]))
withMacros[expr_][nb_][string_String]:=withMacros[expr][nb][{string}];

settingStrings={"StoreFrequency", "Evaluator"};

initializeSettings[frec_, eval_][nb_]:=withMacros[
"StoreFrequency"=frec;
"Evaluator"=eval;
][nb][settingStrings]



getCellID[]:="CellID"/.Developer`CellInformation[EvaluationNotebook[]]

ClearAll[appendEvents];
appendEvents[evList_,(Rule|RuleDelayed)[lhs_,rhs_]]/;MemberQ[evList, (Rule|RuleDelayed)[lhs,_]]:=evList/.(Rule|RuleDelayed)[lhs,sth_]:>(lhs:>(sth;rhs));
appendEvents[evList_, ev_]:=Append[evList, ev]

checkCreateUndoNotebook[]:=Refresh[
With[{nb:=CurrentValue[EvaluationNotebook[], {TaggingRules,"undoNotebook"}]},
If[(* if the undo notebook is not created or open *)!TrueQ@NotebookTools`NotebookOpenQ@nb,
nb=CreateDocument[{}, CellGrouping->Manual,Visible->False,ShowSelection->False,
WindowTitle->FileBaseName[WindowTitle/.AbsoluteOptions[EvaluationNotebook[], WindowTitle]]<>"_undoData"];
]], None]

nbdeRules={Dynamic[rhs_]:>Dynamic[rhs; Needs["Undo`"];Undo`Private`checkCreateUndoNotebook[]], Null:>Dynamic[Needs["Undo`"];Undo`Private`checkCreateUndoNotebook[]],rhs_:>Dynamic[rhs;Needs["Undo`"];Undo`Private`checkCreateUndoNotebook[]]};

parseNBEOptions[nb_]:=ReleaseHold@Replace[Hold[NotebookDynamicExpression]/.Options[nb, NotebookDynamicExpression], nbdeRules, {1}]

appendNDE[nb_NotebookObject]:=SetOptions[nb, parseNBEOptions[nb]/.newNDE_:>NotebookDynamicExpression:>newNDE]

appendNEA[nb_NotebookObject]:=(* appends the notebook event action WindowClose *)
appendEvents[Replace[NotebookEventActions/.Options[nb, NotebookEventActions], None:>{}], "WindowClose":>NotebookClose@CurrentValue[EvaluationNotebook[], {TaggingRules,"undoNotebook"}]]/.newOpts_:>SetOptions[nb, NotebookEventActions:>newOpts]

Options[SetUpUndo]={"Method"->"FrontEnd", "StoreFrequency"->15};
SetUpUndo[nb_NotebookObject,styles_List, OptionsPattern[]]:=SetOptions[initializeNotebook[nb,{OptionValue["StoreFrequency"], OptionValue["Method"]}];nb, StyleDefinitions->Fold[
#1/.{s_String:>Notebook[{Cell[StyleData[StyleDefinitions->s]],Cell[StyleData[#2],cellEventActions
]}],
Notebook[{bef___,Cell[StyleData[#2],cellRest___],aft___}, rest___]:>Notebook[{bef,Cell[StyleData[#2], cellRest, cellEventActions],aft}, rest],
Notebook[{all___}, rest___]:>Notebook[{all, Cell[StyleData[#2], cellEventActions]},rest]
}&,
FullOptions[nb, StyleDefinitions],styles]];

SetUpUndo[nb_,styles_, op:OptionsPattern[]]:=SetUpUndo[nb, {styles},op];

initializeNotebook[nb_NotebookObject, {frec_, eval_}]:=Through@{appendNDE, appendNEA, initializeSettings[frec, eval],
SetOptions[#, System`CreateCellID->True]&}[
nb];

cellEventActions=CellEventActions:>{PassEventsDown -> True, "KeyDown" :> If[CurrentValue["EventKey"]=!=None,undoSave[CurrentValue[CellID]]],"MouseClicked":>If[CurrentValue["AltKey"]&&!NotebookTools`HorizontalInsertionPointQ[EvaluationNotebook[]], undoLoad@CurrentValue@CellID]};

undoFESave[id_Integer]:=undoFESave[ToString@id, CurrentValue[EvaluationNotebook[], {TaggingRules,"undoNotebook"}]];

undoFESave[id_, undoNb_]:=(If[
NotebookFind[undoNb, ToString@id<>"End", All, CellTags]===$Failed,
createCellIDZone[id][undoNb],
SelectionMove[undoNb, Before, Cell]];
copyCell[EvaluationNotebook[], undoNb]);

createCellIDZone[id_][nb_NotebookObject]:=(ToString@id/.idStr_:>{idStr<>"Start", idStr<>"End"}/.{stag_, etag_}:>FrontEndExecute[{
FrontEnd`SelectionMove[nb,After,CellGroup],
FrontEnd`NotebookWrite[nb,Cell@CellGroupData[{Cell["Start", CellTags->stag],Cell["End", CellTags->etag]}]],
FrontEnd`SelectionMove[nb, Previous, Cell],
FrontEnd`SelectionMove[nb, Before, Cell]}];
nb)

idCounter[_]=0;

ClearAll[doSthWithWholeCellSelected];
doSthWithWholeCellSelected[sth_,what_:Cell][nb_]:=With[{val=CurrentValue[nb, ShowSelection],
cursor=("CursorPosition"/.Developer`CellInformation[nb])[[1,1]]},
FrontEndExecute[{
FrontEnd`SetOptions[nb, ShowSelection->False],
FrontEnd`SelectionMove[nb, All, what,AutoScroll->False]}];
With[{cell=sth[nb]},
{FrontEnd`SelectionMove[nb, Before, CellContents, AutoScroll->False],
FrontEnd`SelectionMove[nb, Next, Character, cursor, AutoScroll->False],
FrontEnd`SetOptions[nb, ShowSelection->val]}/.l:{a_, b_, c_}:>If[cursor>0,FrontEndExecute@l,FrontEndExecute[{a,c}]];
cell
]
]

getCurrentCell[nb_NotebookObject]:=doSthWithWholeCellSelected[NotebookRead][nb];

copyCell[nbFrom_, nbTo_]:=NotebookWrite[nbTo, getCurrentCell[nbFrom]]

replaceCurrentCellContents[nb_NotebookObject, newCell_]:=doSthWithWholeCellSelected[NotebookWrite[#, newCell]&, CellContents][nb];

undoFELoad[id_Integer]:=undoFELoad[ToString@id,CurrentValue[EvaluationNotebook[], {TaggingRules, "undoNotebook"}]];

undoFELoad[id_String, nb_]:=If[(* if it is found *)
NotebookFind[nb, ToString@id<>"End", All, CellTags]=!=$Failed,
SelectionMove[nb, Previous, Cell];
replaceCurrentCellContents[EvaluationNotebook[],If[CurrentValue[NotebookSelection@nb, CellTags]=!=ToString@id<>"Start",(NotebookDelete[nb];#)&@NotebookRead[nb],SelectionMove[nb, After, Cell];""]]]

withMacros[
undoSave[id_Integer]:=If[Mod[ ++ idCounter[id],
        "StoreFrequency"] === 0,
      Switch["Evaluator","Kernel",undoKernelSave,"FrontEnd", undoFESave][id]];
undoLoad:=Switch["Evaluator","Kernel",undoKernelLoad,"FrontEnd", undoFELoad];
][Unevaluated@EvaluationNotebook[]][settingStrings]

undoData[_,_Integer]={};

undoKernelSave[ID_Integer]:=undoKernelSave[ID, EvaluationNotebook[]];
undoKernelSave[ID_, nb_]:=undoData[EvaluationNotebook[], ID]={getCurrentCell[nb],undoData[EvaluationNotebook[], ID]};

undoKernelLoad[ID_Integer]:=undoKernelLoad[ID, EvaluationNotebook[]];
undoKernelLoad[ID_, nb_]:=With[{data=undoData[nb,ID]},
(replaceCurrentCellContents[nb,First@data ];
undoData[nb, ID]=Last@data;
)/;data=!={}]

End[];
EndPackage[];