Group.cpp

#include "Group.h"
#include "Record.h"
#include "HAL/FileManager.h"
#include "Serialization/MemoryReader.h"
#include "Serialization/MemoryWriter.h"

// GRUP type identifier as a 4-byte constant ('GRUP').
static const uint32 GroupType = 'G' | ('R' << 8) | ('U' << 16) | ('P' << 24);

// Returns the header length to skip based on game-specific format (12 for Oblivion, 16 for others).
static int32 GetHeaderLengthToSkip(EGameID GameID)
{
    return (GameID == EGameID::Oblivion) ? 12 : 16;
}

// Internal function to read a group from an archive, handling recursive subgroups.
// Parameters: OutGroup (output group), Archive (data source), GameID (game format), bSkipSubrecords (skip record subrecords), Depth (recursion depth).
// Returns true on success, false on failure (e.g., invalid type, archive error).
static bool ReadGroupInternal(FGroup& OutGroup, FArchive& Archive, EGameID GameID, bool bSkipSubrecords, int32 Depth = 0)
{
    // Limit recursion depth to prevent stack overflow on malformed files.
    const int32 MaxDepth = 100;
    if (Depth >= MaxDepth)
    {
        UE_LOG(LogTemp, Warning, TEXT("Maximum group recursion depth reached."));
        return false;
    }

    // Clear existing FormIDs and initialize a set for deduplication.
    OutGroup.FormIds.Empty();
    OutGroup.Label = TEXT(""); // Initialize label to empty.
    TSet<FFormID> FormIdSet;

    // Get game-specific header length to skip (label and other fields).
    int32 HeaderLengthToSkip = GetHeaderLengthToSkip(GameID);

    // Read the 4-byte type identifier.
    uint32 Type = 0;
    Archive.Serialize(&Type, sizeof(uint32));
    if (Type != GroupType)
    {
        UE_LOG(LogTemp, Warning, TEXT("Invalid group type: %u"), Type);
        return false;
    }

    // Read the total group size, including header and contents.
    uint32 GroupSize = 0;
    Archive.Serialize(&GroupSize, sizeof(uint32));
    if (Archive.IsError())
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to read group size."));
        return false;
    }

    // Read the group label (4 bytes, interpreted as string or integer based on game).
    TArray<uint8> LabelBytes;
    LabelBytes.SetNumUninitialized(4);
    Archive.Serialize(LabelBytes.GetData(), 4);
    if (Archive.IsError())
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to read group label."));
        return false;
    }
    OutGroup.Label = FString(4, reinterpret_cast<TCHAR*>(LabelBytes.GetData()));

    // Skip remaining header fields (e.g., stamp, version).
    Archive.Seek(Archive.Tell() + (HeaderLengthToSkip - 4));
    if (Archive.IsError())
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to skip group header."));
        return false;
    }

    // Calculate the size of records/subgroups to process.
    int32 TotalRecordsSize = GroupSize - 8 - HeaderLengthToSkip;
    int64 StartPos = Archive.Tell();

    // Read records or subgroups until the size is consumed or an error occurs.
    while (!Archive.IsError() && (Archive.Tell() - StartPos) < TotalRecordsSize)
    {
        int64 CurrentPos = Archive.Tell();
        uint32 NextType = 0;
        Archive.Serialize(&NextType, sizeof(uint32));
        Archive.Seek(CurrentPos); // Rewind to process the item.

        if (NextType == GroupType)
        {
            // Recursively read a subgroup.
            FGroup SubGroup;
            if (!ReadGroupInternal(SubGroup, Archive, GameID, bSkipSubrecords, Depth + 1))
            {
                UE_LOG(LogTemp, Warning, TEXT("Failed to read subgroup at position %lld."), CurrentPos);
                return false;
            }
            for (const FFormID& FormId : SubGroup.FormIds)
            {
                FormIdSet.Add(FormId); // Deduplicate FormIDs.
            }
        }
        else
        {
            // Read a record and extract its FormID.
            FRecord Record;
            if (!URecordHelper::ReadRecord(Record, Archive, GameID, bSkipSubrecords))
            {
                UE_LOG(LogTemp, Warning, TEXT("Failed to read record at position %lld."), CurrentPos);
                return false;
            }
            FormIdSet.Add(Record.FormId);
        }
    }

    // Convert deduplicated FormIDs to array.
    OutGroup.FormIds = FormIdSet.Array();
    return !Archive.IsError();
}

// Internal function to write a group to an archive, handling header and contents.
// Parameters: Group (input group), Archive (output destination), GameID (game format).
// Returns true on success, false on failure.
static bool WriteGroupInternal(const FGroup& Group, FArchive& Archive, EGameID GameID)
{
    // Write GRUP type identifier.
    uint32 Type = GroupType;
    Archive.Serialize(&Type, sizeof(uint32));
    if (Archive.IsError())
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to write group type."));
        return false;
    }

    // Placeholder for group size (to be updated later).
    uint32 GroupSize = 0;
    int64 SizePos = Archive.Tell();
    Archive.Serialize(&GroupSize, sizeof(uint32));
    if (Archive.IsError())
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to write group size placeholder."));
        return false;
    }

    // Write group label (4 bytes).
    TArray<uint8> LabelBytes;
    LabelBytes.SetNumZeroed(4);
    for (int32 i = 0; i < 4 && i < Group.Label.Len(); ++i)
    {
        LabelBytes[i] = static_cast<uint8>(Group.Label[i]);
    }
    Archive.Serialize(LabelBytes.GetData(), 4);
    if (Archive.IsError())
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to write group label."));
        return false;
    }

    // Write remaining header fields (e.g., stamp, version).
    int32 HeaderLengthToSkip = GetHeaderLengthToSkip(GameID);
    TArray<uint8> HeaderPadding;
    HeaderPadding.SetNumZeroed(HeaderLengthToSkip - 4);
    Archive.Serialize(HeaderPadding.GetData(), HeaderPadding.Num());
    if (Archive.IsError())
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to write group header padding."));
        return false;
    }

    // Write minimal records for each FormID (FGroup lacks full record data).
    for (const FFormID& FormId : Group.FormIds)
    {
        // Write a minimal record header (type "NULL", size 12/16, FormID).
        TArray<uint8> TypeBytes;
        TypeBytes.SetNumZeroed(4);
        FString NullType = TEXT("NULL");
        for (int32 i = 0; i < 4 && i < NullType.Len(); ++i)
        {
            TypeBytes[i] = static_cast<uint8>(NullType[i]);
        }
        Archive.Serialize(TypeBytes.GetData(), 4);
        if (Archive.IsError())
        {
            UE_LOG(LogTemp, Warning, TEXT("Failed to write record type for FormID."));
            return false;
        }

        // Write record size (minimal: flags + FormID + padding).
        uint32 RecordSize = (GameID == EGameID::Oblivion) ? 12 : 16;
        Archive.Serialize(&RecordSize, sizeof(uint32));
        if (Archive.IsError())
        {
            UE_LOG(LogTemp, Warning, TEXT("Failed to write record size."));
            return false;
        }

        // Write flags (0 for minimal record).
        int32 Flags = 0;
        Archive.Serialize(&Flags, sizeof(int32));
        if (Archive.IsError())
        {
            UE_LOG(LogTemp, Warning, TEXT("Failed to write record flags."));
            return false;
        }

        // Write FormID.
        int32 FormIdValue = FormId.Value;
        Archive.Serialize(&FormIdValue, sizeof(int32));
        if (Archive.IsError())
        {
            UE_LOG(LogTemp, Warning, TEXT("Failed to write FormID."));
            return false;
        }

        // Write padding (4 for Oblivion, 8 for others).
        int32 PaddingLength = (GameID == EGameID::Oblivion) ? 4 : 8;
        TArray<uint8> Padding;
        Padding.SetNumZeroed(PaddingLength);
        Archive.Serialize(Padding.GetData(), PaddingLength);
        if (Archive.IsError())
        {
            UE_LOG(LogTemp, Warning, TEXT("Failed to write record padding."));
            return false;
        }
    }

    // Update group size.
    int64 EndPos = Archive.Tell();
    GroupSize = static_cast<uint32>(EndPos - SizePos);
    Archive.Seek(SizePos);
    Archive.Serialize(&GroupSize, sizeof(uint32));
    Archive.Seek(EndPos);
    if (Archive.IsError())
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to update group size."));
        return false;
    }

    return true;
}

// Reads a group from a file, creating an archive and delegating to ReadGroupInternal.
bool UGroupHelper::ReadGroupFromFile(FGroup& OutGroup, const FString& FilePath, EGameID GameID, bool bSkipSubrecords)
{
    TUniquePtr<FArchive> FileArchive(IFileManager::Get().CreateFileReader(*FilePath));
    if (!FileArchive)
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to open file: %s"), *FilePath);
        return false;
    }
    return ReadGroupInternal(OutGroup, *FileArchive, GameID, bSkipSubrecords);
}

// Reads a group from a byte array, creating a memory reader and delegating to ReadGroupInternal.
bool UGroupHelper::ReadGroupFromBytes(FGroup& OutGroup, const TArray<uint8>& ByteArray, EGameID GameID, bool bSkipSubrecords)
{
    FMemoryReader MemoryArchive(ByteArray, true);
    return ReadGroupInternal(OutGroup, MemoryArchive, GameID, bSkipSubrecords);
}

// Writes a group to a file, creating an archive and delegating to WriteGroupInternal.
bool UGroupHelper::WriteGroupToFile(const FGroup& Group, const FString& FilePath, EGameID GameID)
{
    TUniquePtr<FArchive> FileArchive(IFileManager::Get().CreateFileWriter(*FilePath));
    if (!FileArchive)
    {
        UE_LOG(LogTemp, Warning, TEXT("Failed to create file: %s"), *FilePath);
        return false;
    }
    return WriteGroupInternal(Group, *FileArchive, GameID);
}

// Serializes a group to a byte array, creating a memory writer.
bool UGroupHelper::WriteGroupToBytes(const FGroup& Group, TArray<uint8>& OutByteArray, EGameID GameID)
{
    OutByteArray.Empty();
    FMemoryWriter MemoryArchive(OutByteArray, true);
    return WriteGroupInternal(Group, MemoryArchive, GameID);
}

// Returns the FormIDs in the group.
TArray<FFormID> UGroupHelper::GetFormIds(const FGroup& Group)
{
    return Group.FormIds;
}

// Returns the number of FormIDs in the group.
int32 UGroupHelper::GetFormIdCount(const FGroup& Group)
{
    return Group.FormIds.Num();
}

// Checks if a specific FormID exists in the group.
bool UGroupHelper::HasFormId(const FGroup& Group, const FFormID& FormId)
{
    return Group.FormIds.Contains(FormId);
}

// Returns the group label.
FString UGroupHelper::GetLabel(const FGroup& Group)
{
    return Group.Label;
}

// Adds a FormID to the group.
void UGroupHelper::AddFormId(FGroup& Group, const FFormID& FormId)
{
    Group.FormIds.AddUnique(FormId); // Ensure no duplicates.
}

// Removes a FormID from the group.
bool UGroupHelper::RemoveFormId(FGroup& Group, const FFormID& FormId)
{
    return Group.FormIds.Remove(FormId) > 0;
}

// Clears all FormIDs from the group.
void UGroupHelper::ClearFormIds(FGroup& Group)
{
    Group.FormIds.Empty();
}

// Validates the group (non-empty FormIDs, non-empty label).
bool UGroupHelper::IsValidGroup(const FGroup& Group, EGameID GameID)
{
    if (Group.FormIds.Num() == 0)
    {
        return false;
    }
    if (Group.Label.IsEmpty())
    {
        return false;
    }
    // For non-Morrowind games, ensure FormIDs are non-zero.
    if (GameID != EGameID::Morrowind)
    {
        for (const FFormID& FormId : Group.FormIds)
        {
            if (FormId.Value == 0)
            {
                return false;
            }
        }
    }
    return true;
}