// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "symbol_resolver.h"
#include "trace_model.h"

#include <zencore/sharedbuffer.h>

#include <cstdint>
#include <filesystem>
#include <memory>
#include <optional>
#include <string_view>

ZEN_THIRD_PARTY_INCLUDES_START
#include <EASTL/hash_map.h>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen::trace_detail {

// ---------------------------------------------------------------------------
// File format constants
// ---------------------------------------------------------------------------

static constexpr uint32_t kCacheMagic	= 0x005A4355;  // "UCZ\0"
static constexpr uint32_t kCacheVersion = 1;

enum class CacheSectionId : uint32_t
{
	StringTable = 0,
	Metadata	= 1,
	Memory		= 2,
	Callstacks	= 3,
	Symbols		= 4,
	Count
};

// ---------------------------------------------------------------------------
// On-disk header structures (naturally aligned, no packing)
//
// Fields are ordered so natural alignment matches across compilers without
// needing pragma pack. static_asserts at the bottom of this block pin the
// layout so a reordering or added field cannot silently break cached files.
// ---------------------------------------------------------------------------

struct CacheFileHeader
{
	uint32_t Magic;
	uint32_t Version;
	uint64_t SourceFileSize;
	int64_t	 SourceModTimeNs;  // last_write_time as nanoseconds since epoch
	uint64_t Reserved;
};

struct SectionDirectoryEntry
{
	uint32_t SectionId;
	uint32_t Reserved;
	uint64_t FileOffset;	  // byte offset from start of file
	uint64_t CompressedSize;  // size of the CompressedBuffer blob on disk
};

// ---------------------------------------------------------------------------
// POD types for memory-mappable section content
// ---------------------------------------------------------------------------

struct MetadataPod
{
	uint64_t FileSize;
	uint64_t TotalEvents;
	uint64_t ParseTimeMs;
	uint32_t TraceStartUs;
	uint32_t TraceEndUs;
	// SessionInfo string indices
	uint32_t SessionPlatform;
	uint32_t SessionAppName;
	uint32_t SessionProjectName;
	uint32_t SessionCommandLine;
	uint32_t SessionBranch;
	uint32_t SessionBuildVersion;
	uint32_t SessionChangelist;
	uint8_t	 SessionConfigType;
	uint8_t	 SessionHasSession;
	uint8_t	 Padding[2];
};

struct ThreadInfoPod
{
	uint32_t ThreadId;
	uint32_t Name;		 // string index
	uint32_t GroupName;	 // string index
	uint32_t SystemId;
	int32_t	 SortHint;
	uint32_t Pad;
};

struct ChannelInfoPod
{
	uint32_t Name;	// string index
	uint8_t	 Enabled;
	uint8_t	 ReadOnly;
	uint8_t	 Pad[2];
};

struct ModuleInfoPod
{
	uint32_t Name;		// string index
	uint32_t FullPath;	// string index
	uint64_t Base;
	uint32_t Size;
	uint32_t ImageIdSize;	 // byte count in the ImageId blob area
	uint32_t ImageIdOffset;	 // byte offset into the ImageId blob area
	uint32_t Pad;
};

struct EventTypeCountPod
{
	uint32_t Name;	// string index
	uint32_t Pad;
	uint64_t Count;
};

struct CpuScopeStatPod
{
	uint32_t Name;	// string index
	uint32_t MinUs;
	uint32_t MaxUs;
	uint32_t Pad;
	uint64_t Count;
	double	 MeanUs;
	double	 StdDevUs;
};

struct AllocSummaryPod
{
	uint8_t	 HasMemoryData;
	uint8_t	 Pad0[3];
	uint32_t PeakTimeUs;
	uint32_t LiveAllocations;
	uint32_t Pad1;
	uint64_t TotalAllocs;
	uint64_t TotalFrees;
	uint64_t TotalReallocAllocs;
	uint64_t TotalReallocFrees;
	int64_t	 PeakBytes;
	int64_t	 EndBytes;
};

struct HeapInfoPod
{
	uint32_t Id;
	uint32_t ParentId;
	uint16_t Flags;
	uint16_t Pad0;
	uint32_t Name;	// string index
};

struct HeapStatPod
{
	uint32_t HeapId;
	uint32_t Pad;
	int64_t	 CurrentBytes;
	int64_t	 PeakBytes;
	uint64_t AllocCount;
	uint64_t FreeCount;
};

struct CallstackAllocStatPod
{
	uint32_t CallstackId;
	uint32_t LiveCount;
	int64_t	 LiveBytes;
	uint32_t ThreadIdCount;
	uint32_t ThreadIds[4];
	uint32_t Pad;
	uint32_t Pad2;
};

struct CallstackChurnStatPod
{
	uint32_t CallstackId;
	uint32_t Pad;
	uint64_t ChurnAllocs;
	uint64_t ChurnBytes;
	uint64_t TotalAllocs;
	uint64_t TotalBytes;
	double	 MeanDistance;
};

struct CallstackHeaderPod
{
	uint32_t Id;
	uint32_t FrameCount;
	uint32_t FrameOffset;  // index into the frames array
	uint32_t Pad;
};

struct ResolvedFramePod
{
	uint64_t Address;
	uint32_t ModuleIndex;
	uint32_t Pad;
	uint64_t Offset;
};

struct SymbolEntryPod
{
	uint64_t Address;
	uint32_t StringIdx;	 // index into the string table
	uint32_t Pad;
};

// Pin the on-disk layout. Any change here is a cache format change and must
// bump kCacheVersion.
static_assert(sizeof(CacheFileHeader) == 32);
static_assert(sizeof(SectionDirectoryEntry) == 24);
static_assert(sizeof(MetadataPod) == 64);
static_assert(sizeof(ThreadInfoPod) == 24);
static_assert(sizeof(ChannelInfoPod) == 8);
static_assert(sizeof(ModuleInfoPod) == 32);
static_assert(sizeof(EventTypeCountPod) == 16);
static_assert(sizeof(CpuScopeStatPod) == 40);
static_assert(sizeof(AllocSummaryPod) == 64);
static_assert(sizeof(HeapInfoPod) == 16);
static_assert(sizeof(HeapStatPod) == 40);
static_assert(sizeof(CallstackAllocStatPod) == 48);
static_assert(sizeof(CallstackChurnStatPod) == 48);
static_assert(sizeof(CallstackHeaderPod) == 16);
static_assert(sizeof(ResolvedFramePod) == 24);
static_assert(sizeof(SymbolEntryPod) == 16);

// ---------------------------------------------------------------------------
// Cache read / write API
// ---------------------------------------------------------------------------

struct CachedAnalysis
{
	TraceModel						Model;
	std::unique_ptr<SymbolResolver> Symbols;
};

// Try to load a cached analysis from the .ucache_z file next to a .utrace.
// Returns nullopt on any failure (missing, stale, corrupt, version mismatch).
std::optional<CachedAnalysis> TryLoadAnalyzeCache(const std::filesystem::path& CachePath, const std::filesystem::path& SourcePath);

// Write the analysis cache for future reuse.
void WriteAnalyzeCache(const std::filesystem::path&					 CachePath,
					   const std::filesystem::path&					 SourcePath,
					   const TraceModel&							 Model,
					   const eastl::hash_map<uint64_t, std::string>& ResolvedSymbols);

}  // namespace zen::trace_detail