#if SPECIAL
#include "main.h"
#endif

struct history_info {
    uint16 ActionCount_Total;
    uint16 ActionCount_EndOfPlayhead;
    uint16 ActionCount_Current;

    uint64 ActionOffset_Total;
    uint64 ActionOffset_EndOfPlayhead;
    uint64 ActionOffset_Current;

    uint64 EntrySize_Current;
};

static uint64
History_GetActionSize(history_entry_list *History, int i)
{
    history_action *Action = &History->Action[i];
    if (Action->Type == action_type_swap || Action->Type == action_type_swap_bitmap)
        return Action->Size;
    return 0;
}

// Returns information on the undo tree based on three points of data: the
// total amount of entries, the location of the playhead, and whatever location
// SampleIndex is set to, which should be either one minus playhead or the playhead.

// I wrote this with the intent of precomputing things like the tree's total
// size and the size of the current action, because calculating them in the
// Action_Redo/Undo calls got extremely confusing.

static history_info
History_GetTreeInfo(history_entry_list *History, uint16 SampleIndex)
{
    history_info Info = {};

    for (int i = 0; i < History->NumberOfEntries; i++)
    {
        history_entry *Entry = &History->Entry[i];
        Info.ActionCount_Total += Entry->NumberOfActions;
        if (i < History->EntryPlayhead)
            Info.ActionCount_EndOfPlayhead += Entry->NumberOfActions;
        if (i < SampleIndex)
            Info.ActionCount_Current += Entry->NumberOfActions;
    }

    for (int i = 0; i < Info.ActionCount_Total; i++)
    {
        uint64 Size = History_GetActionSize(History, i);
        Info.ActionOffset_Total += Size;
        if (i < Info.ActionCount_EndOfPlayhead)
            Info.ActionOffset_EndOfPlayhead += Size;

        if (i < Info.ActionCount_Current)
            Info.ActionOffset_Current += Size;
        else if (i < Info.ActionCount_EndOfPlayhead)  // Only increment the current size if i is between
            Info.EntrySize_Current += Size;           // the current count and playhead count!
    }

    return Info;
}

// This function works exactly the same whether we're in undo/redo, pretty neat.
void History_Action_DoSwapBitmap(memory *Memory, history_action *ActionPointer, uint8 *Address_HistoryTree_Start, uint8 *Address_HistoryTree_End, uint8 *Address_Data)
{
    history_action Action = *ActionPointer;

    uint64 FullSize = Action.ShiftAmount;
    uint64 CompressedSize = Action.Size;

    uint32 BytesForCompressedPart = (uint32)100 * 1024 * 1024;
    void *UncompressedTempBuffer = Memory_PushScratch(Memory, FullSize);
    void *CompressedTempBuffer = Memory_PushScratch(Memory, BytesForCompressedPart);

    // First we decompress and swap the saved pixels with the newer ones...
    Data_Decompress(Memory, Address_HistoryTree_Start, CompressedSize, UncompressedTempBuffer, FullSize);
    Bitmap_SwapData((uint8 *)UncompressedTempBuffer, (uint8 *)Address_Data, FullSize, Action.Direction);

    // Then we recompress the old pixels and evaluate its size, and since the size can be different...
     uint64 NewSize = Data_Compress(Memory, UncompressedTempBuffer, FullSize, CompressedTempBuffer, BytesForCompressedPart, Z_BEST_SPEED);

    // we have to shift the undo stack.
    if (NewSize != CompressedSize) {
        int16 ShiftDirection = (NewSize > CompressedSize) ? 1 : -1;
        uint64 ShiftAmount = (NewSize > CompressedSize) ? NewSize - CompressedSize : CompressedSize - NewSize;
        Arbitrary_ShiftData(Address_HistoryTree_Start + CompressedSize, Address_HistoryTree_End, ShiftAmount, ShiftDirection);
    }
    Memory_Copy(Address_HistoryTree_Start, (uint8 *)CompressedTempBuffer, NewSize);

    ActionPointer->Size = NewSize;
    Memory_PopScratch(Memory, BytesForCompressedPart);
    Memory_PopScratch(Memory, FullSize);
}

void History_Action_Undo(memory *Memory, history_info Info, history_action *ActionPointer, uint64 Action_Offset)
{
    history_entry_list *History = &Memory->History;

    // Only swap_bitmap should touch the data!
    history_action Action = *ActionPointer;

    uint8 *Address_HistoryTree_Start = (uint8 *)Memory_AddressAtOffset(Memory, P_UndoBuffer, Action_Offset);
    uint8 *Address_HistoryTree_End = (uint8 *)Memory_AddressAtOffset(Memory, P_UndoBuffer, Info.ActionOffset_Total);
    uint8 *Address_Data = (uint8 *)Memory_AddressAtOffset(Memory, Action.TableName, Action.ByteOffset);

    if (Action.Type == action_type_swap) {
        Arbitrary_SwapData(Memory, Address_HistoryTree_Start, Address_Data, Action.Size);
    } else if (Action.Type == action_type_swap_bitmap) {
        History_Action_DoSwapBitmap(Memory, ActionPointer, Address_HistoryTree_Start, Address_HistoryTree_End, Address_Data);
    } else if (Action.Type == action_type_shift) {
        // In order to shift back we have to start the address at where the
        // shifted chunk is _now_, not when we called the function.
        uint64 NewByteOffset = (Action.Direction > 0) ? Action.ByteOffset + Action.ShiftAmount : Action.ByteOffset - Action.ShiftAmount;
        void *Address_Start = Memory_AddressAtOffset(Memory, Action.TableName, NewByteOffset);
        uint8 *Address_End = (uint8 *)Address_Start + Action.Size;
        // Direction also needs to be reversed.
        int16 Direction = Action.Direction * -1;
        Arbitrary_ShiftData((uint8 *)Address_Start, Address_End, Action.ShiftAmount, Direction);
    } else {
        Assert(0);
    }
}

void History_Action_Redo(memory *Memory, history_info Info, history_action *ActionPointer, uint64 Action_Offset) {

    history_entry_list *History = &Memory->History;

    history_action Action = *ActionPointer;

    uint8 *Address_HistoryTree_Start = (uint8 *)Memory_AddressAtOffset(Memory, P_UndoBuffer, Action_Offset);
    uint8 *Address_HistoryTree_End = (uint8 *)Memory_AddressAtOffset(Memory, P_UndoBuffer, Info.ActionOffset_Total);
    uint8 *Address_Data = (uint8 *)Memory_AddressAtOffset(Memory, Action.TableName, Action.ByteOffset);

    if (Action.Type == action_type_swap) {
        Arbitrary_SwapData(Memory, Address_HistoryTree_Start, Address_Data, Action.Size);
    } else if (Action.Type == action_type_swap_bitmap) {
        History_Action_DoSwapBitmap(Memory, ActionPointer, Address_HistoryTree_Start, Address_HistoryTree_End, Address_Data);
    } else if (Action.Type == action_type_shift) {
        void *Address_Start = Memory_AddressAtOffset(Memory, Action.TableName, Action.ByteOffset);
        uint8 *Address_End = (uint8 *)Address_Start + Action.Size;
        int16 Direction = Action.Direction;
        Arbitrary_ShiftData((uint8 *)Address_Start, Address_End, Action.ShiftAmount, Direction);
    } else {
        Assert(0);
    }
}

void History_Undo(memory *Memory) {
    history_entry_list *History = &Memory->History;
    if (History->EntryPlayhead == 0) return;

    // We want Current to represent the beginning of the current entry, so we
    // subtract one from the playhead.
    history_info Info = History_GetTreeInfo(History, History->EntryPlayhead - 1);
    history_entry *Entry = &History->Entry[History->EntryPlayhead - 1];

    // This puts us at the end of the current entry's offset.
    uint64 ActionOffset_Stepback = Info.ActionOffset_Current + Info.EntrySize_Current;

    for (int i = Entry->NumberOfActions - 1; i >= 0; i--) {
        history_action *Action = &History->Action[Info.ActionCount_Current + i];

        // We step backwards only if the action is currently storing data.
        if (Action->Type != action_type_shift)
            ActionOffset_Stepback -= Action->Size;

        History_Action_Undo(Memory, Info, Action, ActionOffset_Stepback);
    }

    History->EntryPlayhead--;

}

void History_Redo(memory *Memory) {
    history_entry_list *History = &Memory->History;
    if (History->EntryPlayhead == History->NumberOfEntries) return;

    // NOTE(fox): The third part of this function for recording the "current"
    // action's size is only necessary for the undo call, so it's not correct on the redo.
    history_info Info = History_GetTreeInfo(History, History->EntryPlayhead);
    history_entry *Entry = &History->Entry[History->EntryPlayhead];

    History->EntryPlayhead++;

    uint64 ActionOffset_Stepforward = Info.ActionOffset_Current;

    for (int i = 0; i < Entry->NumberOfActions; i++) {
        history_action *Action = &History->Action[Info.ActionCount_Current + i];

        History_Action_Redo(Memory, Info, Action, ActionOffset_Stepforward);

        // We step forwards only if the action is currently storing data.
        if (Action->Type != action_type_shift)
            ActionOffset_Stepforward += Action->Size;

    }

}

// This is only called when we're certain the action is going to be taken.
void History_Entry_Commit(memory *Memory, char *Name)
{
    history_entry_list *History = &Memory->History;
    history_entry *Entry = &History->Entry[History->EntryPlayhead];
    Entry->Name = Name;
    Entry->NumberOfActions = 0;
    if (History->NumberOfEntries > MAX_HISTORY_ENTRIES) {
        Assert(0);
    }
    // Effectively deletes entries in front if we're beginning out of an undo.
    if (History->NumberOfEntries != History->EntryPlayhead)
        History->NumberOfEntries = History->EntryPlayhead;
    History->NumberOfEntries++;
    History->EntryPlayhead++;
    Memory->IsFileSaved = false;
    Memory->PurgeCache = true;
#if DEBUG
    Assert(Debug.UndoState != 1);   // You forgot to end a History_Entry_Commit()!
    Debug.UndoState = 1;
#endif
}

void History_Entry_End(memory *Memory)
{
    history_entry_list *History = &Memory->History;
#if DEBUG
    Debug.UndoState = 0;
#endif
}

static void
History_Purge(memory *Memory, history_entry_list *History, uint64 ActionCount_Total, uint64 ActionOffset_Total, uint64 PurgeSize, uint16 PurgeCount = 0) 
{
    int Size = 0;
    int EntryIndex = 0;
    int ActionIndex = 0;
    int ActionCount = 0;
    bool32 Test = 1;
    while (Test) {
        history_entry *Entry = &History->Entry[EntryIndex];
        ActionCount += Entry->NumberOfActions;
        if (ActionCount >= 4060)
            int b = 0;
        while (ActionIndex < ActionCount) {
            Size += History_GetActionSize(History, ActionIndex);
            ActionIndex++;
        }
        EntryIndex++;
        Test = (PurgeCount == 0) ? (Size < PurgeSize) : (EntryIndex < PurgeCount);
    }
    int EntryCount = (EntryIndex + 1);
    {
        uint8 *Address_Start = (uint8 *)Memory_AddressAtOffset(Memory, P_UndoBuffer, Size);
        uint8 *Address_End = (uint8 *)Memory_AddressAtOffset(Memory, P_UndoBuffer, ActionOffset_Total);
        Arbitrary_ShiftData(Address_Start, Address_End, Size, -1);
    }
    {
        uint8 *Address_Start = (uint8 *)&History->Entry[EntryCount];
        uint8 *Address_End = (uint8 *)&History->Entry[History->NumberOfEntries];
        Arbitrary_ShiftData(Address_Start, Address_End, sizeof(history_entry) * EntryCount, -1);
    }
    {
        uint8 *Address_Start = (uint8 *)&History->Action[ActionCount];
        uint8 *Address_End = (uint8 *)&History->Action[ActionCount_Total];
        Arbitrary_ShiftData(Address_Start, Address_End, sizeof(history_action) * ActionCount, -1);
    }
    History->NumberOfEntries -= EntryCount;
    History->EntryPlayhead -= EntryCount;
}

// NOTE(fox): Shift is the only action that additionally changes the data after
// the info is put on the tree, since it's always the same function used.

static void History_Action_Add(memory *Memory, history_action ActionData, void *ReferenceData = NULL)
{
    Assert(Debug.UndoState == 1);   // This action wasn't preceded by a History_Undo_Start!
    history_entry_list *History = &Memory->History;
    history_entry *Entry = &History->Entry[History->EntryPlayhead - 1];

    history_info Info = History_GetTreeInfo(History, History->EntryPlayhead - 1);

    if ((Info.ActionOffset_EndOfPlayhead + ActionData.Size) > Memory->Slot[P_UndoBuffer].Size) {
        History_Purge(Memory, History, Info.ActionCount_Total, Info.ActionOffset_Total, ActionData.Size * 4);
        Entry = &History->Entry[History->EntryPlayhead - 1];
        Info = History_GetTreeInfo(History, History->EntryPlayhead - 1);
    }
    if (Info.ActionCount_Total >= MAX_HISTORY_ACTIONS) {
        History_Purge(Memory, History, Info.ActionCount_Total, Info.ActionOffset_Total, ActionData.Size * 4, History->NumberOfEntries / 4);
        Entry = &History->Entry[History->EntryPlayhead - 1];
        Info = History_GetTreeInfo(History, History->EntryPlayhead - 1);
    }

    history_action *Action = &History->Action[Info.ActionCount_Total];
    *Action = ActionData;

    if (Action->Type == action_type_swap)
    {
        void *Address_HistoryTree = Memory_AddressAtOffset(Memory, P_UndoBuffer, Info.ActionOffset_Total);
        void *Address_Data = Memory_AddressAtOffset(Memory, Action->TableName, Action->ByteOffset);
        Memory_Copy((uint8 *)Address_HistoryTree, (uint8 *)Address_Data, Action->Size);
    } else if (Action->Type == action_type_swap_bitmap)
    {
        // The address of data to be stored in the undo tree is not the same as
        // the address to be recorded, so we have to pass it manually.
        void *Address_HistoryTree = Memory_AddressAtOffset(Memory, P_UndoBuffer, Info.ActionOffset_Total);
        void *Address_Data = ReferenceData;

        // Libz/miniz requires passing an amount of available space, and it
        // returns how much space the compressed data actually is.
        // We don't need to alloc any more memory since we can just use the
        // history tree as the destination.
        // ShiftAmount is also being used to save the uncompressed size.
        uint32 BytesForCompressedPart = (uint32)100 * 1024 * 1024;
        Action->Size = Data_Compress(Memory, Address_Data, Action->ShiftAmount, Address_HistoryTree, BytesForCompressedPart, Z_BEST_SPEED);

    } else if (Action->Type == action_type_shift)
    {
        void *Address_Start = Memory_AddressAtOffset(Memory, Action->TableName, Action->ByteOffset);
        uint8 *Address_End = (uint8 *)Address_Start + Action->Size;
        Arbitrary_ShiftData((uint8 *)Address_Start, Address_End, Action->ShiftAmount, Action->Direction);
    }

    Entry->NumberOfActions++;
}

// Helper functions for different tables.

static void History_Action_BitmapPaint(memory *Memory, uint64 Size_Uncompressed,
                                       void *SourceBitmapAddress, void *CachedPaintAddress, int16 BytesPerPixel)
{
    void *Address_Start = Memory->Slot[F_PrincipalBitmaps].Address;
    uint64 ByteOffset = (ptrsize)SourceBitmapAddress - (ptrsize)Address_Start;
    history_action Action = { F_PrincipalBitmaps, action_type_swap_bitmap, 0, ByteOffset, Size_Uncompressed, BytesPerPixel };
    History_Action_Add(Memory, Action, CachedPaintAddress);
}

static void History_Action_Shift(memory *Memory, memory_table_list TableName,
                                 void *Address0, void *Address1, uint64 Amount, int16 Direction)
{
    void *Address_Start = Memory->Slot[TableName].Address;
    uint64 ByteOffset = (ptrsize)Address0 - (ptrsize)Address_Start;
    uint64 Size = (ptrsize)Address1 - (ptrsize)Address0;

    history_action Action = { TableName, action_type_shift, Size, ByteOffset, Amount, Direction };

    History_Action_Add(Memory, Action);
}

static void History_Action_Block_Swap(memory *Memory, memory_table_list TableName, void *Address_Data)
{
    void *Address_Start = Memory->Slot[TableName].Address;
    uint64 Size = Memory->Slot[TableName].Block_ElementSize;
    uint64 ByteOffset = (ptrsize)Address_Data - (ptrsize)Address_Start;
    history_action Action = { TableName, action_type_swap, Size, ByteOffset, 0 };
    History_Action_Add(Memory, Action);
}

// Remember to dereference pointers if taking the sizeof() a variable, or else Size will be 8!

static void History_Action_Swap(memory *Memory, memory_table_list TableName, uint64 Size, void *Address_Data)
{
    void *Address_Start = Memory->Slot[TableName].Address;
    uint64 ByteOffset = (ptrsize)Address_Data - (ptrsize)Address_Start;
    history_action Action = { TableName, action_type_swap, Size, ByteOffset, 0 };
    History_Action_Add(Memory, Action);
}